Personalized Collaborative Learning with Affinity-Based Variance Reduction
Authors: Chenyu Zhang
Navid Azizan
Venue: ICLR 2026 Poster (2026)
Keywords: personalized collaborative learning
multi-agent systems
federated learning
heterogeneity
personalization
stochastic approximation
variance reduction
Primary area: learning theory
TL;DR: We show that collaboration between arbitrarily heterogeneous agents can yield fully personalized solutions with an adaptive, affinity-based speedup.

Abstract: Multi-agent learning faces a fundamental tension: leveraging distributed collaboration without sacrificing the personalization needed for diverse agents. This tension intensifies when aiming for full personalization while adapting to unknown heterogeneity levels—gaining collaborative speedup when agents are similar, without performance degradation when they are different. Embracing the challenge, we propose personalized collaborative learning (PCL), a novel framework for heterogeneous agents to collaboratively learn personalized solutions with seamless adaptivity. Through carefully designed bias correction and importance correction mechanisms, our method AffPCL robustly handles both environment and objective heterogeneity. We prove that AffPCL reduces sample complexity over independent learning by a factor of $\max\\{n^{-1}, \delta\\}$, where $n$ is the number of agents and $\delta\in[0,1]$ measures their heterogeneity. This *affinity-based* acceleration automatically interpolates between the linear speedup of federated learning in homogeneous settings and the baseline of independent learning, without requiring prior knowledge of the system. Our analysis further reveals that an agent may obtain linear speedup even by collaborating with arbitrarily dissimilar agents, unveiling new insights into personalization and collaboration in the high heterogeneity regime.

Published as a conference paper at ICLR 2026
PERSONALIZED COLLABORATIVE LEARNING WITH AFFINITY-BASEE VARIANCE REDUCTION
Chenyu Zhang Massachusetts Institute of Technology zcysxyêmi edd
Navid Azizan
Massachusetts Institute Technology i zan @md edu
ABSTRACT
Multi-agent learning faces a fundamental tension: leveraging distributed collab- oration without sacrificing the personalization needed for diverse agents. This tension intensifies when aiming for full personalization while adapting to unknown heterogeneity levels gaining collaborative speedup when agents are similar, with- out performance degradation when they are different. Embracing the challenge, we propose personalized collaborative learning (PCL), novel framework for het- erogeneous agents to collaboratively learn personalized solutions with seamless adaptivity. Through carefully designed bias correction and importance correction mechanisms, method AffPCL robustly handles both environment and objective heterogeneity. We prove that AffPCL reduces sample complexity over independent learning by factor of max (n-1, 10), where is the number of agents and & [0, measures their heterogeneity. This affinity-basee polates between the linear speedup of federated learning homogeneous settings and the baseline of independent learning, without requiring prior knowledge of the system. Our analysis further reveals that an agent may obtain linear speedup even by collaborating with arbitrarily dissimilar agents, unveiling new insights into personalization and collaboration in the high heterogeneity regime.
acceleration automatically inter-
INTRODUCTION
Heterogeneity is a defining yet formidable characteristic multi-agent systems. When agents differ significantly, their incentives to collaborate diminish, leveraging experience from others can introduce bias and impede their own learning. This challenge intensifies in scenarios where strategic agents seek highly accurate, tailored solutions. Collaborative multi-agent systems commonly adopt federated learning (FL) setup, where agents communicate via a central server to jointly learn a unified solution. However, the presence of heterogeneity, such unified solutions often prove suboptimal or even irrelevant for individual agents. Consequently, effective personalization becomes essential for collaborative learning among heterogeneous agents.
This need is evident in real-world applications: personalized recommendations drive user engagement (Good al. 1999 Anand & Mobasher [2003] Khribi et al. 2008). autonomous transportation must
accommodate local traffic conditions (Huang et al. 2021 You et al. 2024 Jalota et al. [223) robots
need adapt to different operating environments (Richards et al. 2023 lang et al.) (025), diverse patient profiles require tailored treatments (Chen et al. 20222] Tang et al.
2024), and agentic language
models need to adapt to specific user styles and task contexts (Li et al.
2024 Wožniak al. 2024
Bose et al. 2025 Nobari et al. 2025).
These considerations motivate the following multi-agent decision maainng setup.
1. Personalized. Agents are intrinsically heterogeneous, each with arbitrarily distinct environments and objectives, and act strategically to optimize their own goals.
Collaborative. Agents communicate through central server that aggregates information from agents and broadcasts back the aggregated result.
Learning. Agents have no prior knowledge of their systems and interact only with local environ- ments that generate stochastic observations of system parameters.

---

Published as a conference paper at ICLR 2026
Such complex, stochastic, and heterogeneous multi-agent system demands, but also challenges. the design of a personalized collaborative learning algorithm that can (1) find fully personalized solutions for all agents, (2) achieve performance gains through collaboration, (3) and adapt to unknown heterogeneity among agents without prior knowledge, automatically harnessing greater collaboration benefits when agents are similar, and falling back to, while ensuring worse performance than, non-collaborative independent learning when agents are markedly different.
This work reveals that the key to achieving these goals lies in identifying and exploiting affinity, i.e., similarity among agents, Formally, we capture agent heterogeneity through a score ó E [0, 1]. with 5 0 indicating homogeneous agents and larger values of ó indicating greater heterogeneity. For any agent, our method finds its personalized solution with a mean squared error of
O(1-² max (n (n- ;))),
where is the number of samples collected by each agent and n is the number of agents. This finite- sample complexity enjoys federated speedup linear in when agents are similar, while it gracefully reduces to the baseline rate of independent learning O(t when agents are highly heterogeneous, but never worse. In intermediate regimes, affinity-based acceleration manifests.
We summarize our main contributions:
1. We formulate novel multi-agent decision decision-making paradigm of personalized collaborative learning (PCL., encompassing applications and problems in supervised learning, reinforcement learning (Appendix (.6)) and statistical decision-making.
We develop simple yet effective method that realizes the vision of PCL, called AffPCL, which finds fully personalized solutions and adaptively harnesses collaboration benefits when agents are similar while ensuring no worse performance than independent learning when they are highly heterogeneous. Our method robustly handles arbitrary objective and environment heterogeneity through principled personalized bias correction and importance correction mechanisms.
3. We establish finite-sample convergence guarantees for AtfPCL, achieving the rate in (T) and thus demonstrating the desired phenomenon of affinity-bssed variance reduction. This rate adaptively interpolates between the linear speedup of FL and the minimax optimal rate of independent learning. This is the first result that proves efficiency gains for learning fully personalized solutions through collaboration among arbitrarily heterogeneous agents.
4. We further enhance AfIPLL with features including asynchronous importance estimation and agent- specific update schemes. Our agent-specific analysis reveals that an agent may achieve linear speedup even when it is dissimilar to all others, a phenomenon unattainable in prior frameworks.
1.1 RELATED WORK
We focus on the most relevant works in heterogeneous collaborative learning that motivate this study.
Personalization falls short in federated learning. Classical FL methods (McMahan et al. 2017)
aim for a unified solution for all agents without personalization guarantees. The unified objective mitigates heterogeneity; for instance, bias correction in heterogeneous FL. (Karimireddy al. 2020b) Guo et al. 2023 Karimireddy et al.] 2020a; Gao et al. (2022). which prevents local updates from drifting away from the central update direction, averaged across all agents and thus enjoys federated variance reduction (see also Section 2). In contrast, personalization requires adjusting the central update relative to each agent's unique local direction, which precludes federated variance reduction.
The growing literature on (partially) personalized FL highlights the importance of personalization. A common strategy combines global and local models through regularization or mixtures (Li et al.) 2022)) Hanzely & Richtárik 2020 T Dinh et al. 2020 Li et al. [001 Deng et al., 2020), but such methods offer only partial personalization and the trade-ofts may be heuristic. Similarly, clustering-baa based
methods (Sattler et al., 2020; Mansour et al. 2020 Ghosh et al., 2020 Briggs al. 2020 Chai et al., 2020 GGrimeer et al. 2021> do not offer personalization within each cluster and may be sensitive to hyperparameter tuning prior knowledge. In contrast, PCL aims for full personalization and seamless adaptivity, requiring neither prior knowledge of heterogeneity nor hyperparameter tuning.
Slower rates in independent learning. Other personalized learning approaches combine FL and independent learning. A sequential strategy uses FL as a warm start followed by independent fine- tnnin (Fallah al. 2020 Cheng al. (021); while effective in some cases, this approach is
2

---

Published as a conference paper at ICLR 2026
generally rate-suboptimal, the small initialization error through FL diminishes faster than the variance from independent learning, making its change in finite-tima complexity marginal. A parallel approach simultaneously learns shared global component and a personalized local component (Pillutla et al., 2022; Xiong et al. 2025 Liang et al. 2020); this approach requires certain global-local structures, and similarly, the independent learning component dominates the overall complexity, obscuring collaborative speedup. In contrast, PCL imposes no structural assumptions, accommodates arbitrarily heterogeneous agents, and aims for provably faster rates than independent learning.
Curse of heterogeneity in collaborative learning. Closest to our setup, Chayti et al. [202] 1 )
Even et al. (2022) also study full personalization with arbitrarily heterogeneous systems, but with fundamentally different approaches from ours in handling heterogeneity to achieve collaborative variance reduction. First, they selectively collaborate with similar agents, effectively reducing to clustering- hased methods or low heterogeneity regimes, whereas AffPCL enables collaboration among all agents regardless of similarity. With AffPCL, an agent may attain linear speedup even when it's not similar to any other agent (Section 6.3). which is unattainable in their frameworks. Second, achieving optimal speedup in their setting requires either knowledge of objective heterogeneity (Even et al. 2022) access to bias estimation oracle whose variance reduces linearly in the number of agents (Chayti et al. (02T), which is a strong assumption as bias estimation for personalization is inherently agent-specific, and its variance does not reduce with more agents. In contrast, AffPCL requires no prior knowledge or bias estimation oracle, and enjoys affinity-based variance reduction fully adaptively.
1.2 PROBLEM FORMULATION
We consider general multi agent linear system:
(2) where sym (A') 3 + (A")") > O. Each agent aims to find the fixed point zi, of its system with access to only stochastic observations A(st) Rgddd and b' b'(s:) E Rd evaluated its local random state st E S independently sampled from its distinct environment distribution μ E ASS) at time step The stochastic observations are unbiased such that A' A(s)) and b' = (s")
A'x; = B², 1=
Terminology and notation. Our system modeling draws inspiration from various fields, including supervised learning, reinforcement learning, and statistical decision decision-making, where (A, μ)) are commonly referred to as (feature, label, covariate distribution), (function approximation, reward, stationary distribution), and (measurement, response, data distribution), respectively. To appeal to a broader audience and align with our setup, we refer to A as the feature embedding matrix, b as the objective vector, and ju as the environment distribution. As is common in practice, we assume that all agents share the same feature extractor A, but may have different objectives b² and environments j²,, referred to as objective heterogeneity and environment heterogeneiry, respectively.
Throughout the paper, superscript denotes quantities related to agent and superscript 0 denotes the averaged quantity across all agents, i.c., fo = Σi= f for any quantity f. The averaged quantity may be explicitly aggregated by the central server, or it can represent a virtual quantity only used for analysis. We write [m]] == {1,..., n1) For any function i on S, f denotes the expectation of f under the corresponding environment distribution p², i.e., f = E,, E' (s"). For an unknown quantity
f. its estimate learned at time step f is denoted by ft. The default norm is the Fuclidean norm for vectors, operator norm for matrices, and total variation norm for distribution differences. Appendix A contains a complete list of notation.
Roadmap. This paper adopts a progressive approach first develop insights in stylized settings and then incrementally extend to more complex scenarios. We start with a simplified FL setup (Section 2). then gradually introduce personalization (Section (3), adaptivity (Section 4 environment heterogeneity (Section[5.1) and finally arrive at the most general setup (2) in Section (5.2) Several theoretical extensions are discussed in Section 66 and numerical results are presented in Section 7
2 WARM-UP: HETEROGENEOUS FEDERATED LEARNING
We start by reviewing heterogeneous FL, a variant of where agents with distinct objectives collaborate to find unified solution xς satisfying
3

---

Published as a conference paper at ICLR 2026
B [s) This warm-up section assumes homogeneous environment distributions n² = μa for all i [m], and thus can drop the superseript of A. Since sym(A) (A) is positive definite, in a federated stochastic approximation setting, each agent adopts the following fixed-point iteration:
where = E and 8 -m=
2|+1 2 (x)) where gi(xt) A(s))at b(s)),
where dr is the step size and the update direction of is the stochastic residual at time step
To focus on the main ideas, this work considers simplified communication scheme, where agents communicate with a central server at every time step. FL, agents send their local updates to the server, which aggregates them to get the central decision variable 2[41 and broadcasts it back r/-11 1 f 2/+1- The resultant central update rule is then given by
1 n2 Σ
! n E
AG)) 1 b
. =t o0g (xf), where (x
We note that in this FL setting, the local decision variables are always synced with the central one, and thus we have gi (at) gf (xf). Moreover, we can write the central decision variable as the average of the local ones: 2x - =x zy. However, this equivalence becomes obsolete when introduce heterogeneous environments and personalization.
Constants. We define the following constants used throughout. = min, Amin (sym(A')) (A')) ensures strong monotonicity of the fixed-point iteration and controls the convergence rate; an analogous condition in optimization is A-strong convexity A-PL condition of the objective function (Nesterov) 2013). GA = max; sup, (s) Ge = max; sup, || (s)ll, and Gz = max; || upper bound the system parameters. Let a 2 max[G G. represent the seale of the system, which can also be thought of as the variance proxy of the update direction at the solution point, since l|s] [|] S || +|bb(s S G.G. + G, S o; its analogy in optimization the objective function gradient's Lipschitz constant. We then define = σ/À as the condition number of the stochastic system. Without loss of generality, we use as the variance proxy the environment distributions, in the sense that tr Var (f()) E, (s) G7. which holds for any cero-mean
operator f with ess sup [+] Gf.
We have the following convergence guarantee for heterogeneous FI. Proposition With constant step size (At), (4) satisfies
E E - 21 Ō
where Ô suppresses the logarithmic dependence on In
The mean squared error (MSE) of FL vanishes linearly & goes to infinity, with the rate scaled by the
problem scale a and controlled by A. The federated collaboration contributes linear speedup in terms
of the number of agents n. Proposition[]] is tight K, f. and n (Woodworth al. 2020 Karimireddy
et al. 2020b Glasgow et al.) 2022), and serves as a baseline for our subsequent results.
INTRODUCING PERSONALIZATION: PERSONALIZED BIAS CORRECTION
Due to heterogeneity, the unified solution described in Section is generally suboptimal for individual agents, and becomes less relevant as the heterogeneity level grows. More realistically, strategic agents seek personalized solutions:
Áut = 8 i E [m]]
'The unified- solution setting is also related to distributed linear-system solvers Sziizan-Ruba et al. 2019 Velasevic et al. 2023). but those assume that the equations are partitioned across agents.
*Al proot' are deferred to Appendices to G where progressively establish the main result Theorem and cover all the propositions the main
'The t dependence can be removed by using linearly diminishing step size and considering a convex combination of the iterates {ar H as specified in Lemma D.5 This refinement applies to all results in the main text. For brevity, we defer the related discussion to the appendix and omit this remark in subsequent results.
4

---

Published as a conference paper at ICLR 2026
To build intuition, this section makes two simplifications to be relaxed in the next two sections:
agents have the same environment distribution, and the central objective (0(s)) = Σi= b% (s is known to agent upon observing st. With access to the central objective, we propose affinity-aware personalized collaborative learning (AfiPCL), simple yet effective update rule for each agent:
a = = argi where = ( + ( g
gt (x+).
(5)
where the update direction consists of three components:
(( A(5))s) (6)),
- (()),
of (0 A A(s))a - (st)-
Recall that x + E. ox is syneed with the central server. Alternatively, inspired by Section 2 we can replace z with an explicitly maintained central decision variable rf and update it using (4) within the same communication round for computing the central update direction. Both implementations have the same convergence guarantee in current setting, while the latter proves robust to heterogeneous environment distributions, detailed in Section 5.2 See Appendix C. for for further discussion.
Unlike FL., the convergence of AffPCL depends on how similar the objectives of agents are. Definition (Objective heterogeneity). The objective heterogeneity level is defined as
Sohj max sup [lb((s) (s) 8 ()) E [0, 1].
iJj(m) sES
Proposition 2. With constant step size a = In (Mt), satisfies
E||at 21 = Ō(k²t max(n bon)),, Vi [m],
where days S min [1, } is the effective objective heterogeneity level.
The precise definition of the effective heterogeneity level is 6.gj min(1, vōotj where is the stochastic condition number that is trivially bounded by к. We defer the definition of and the discussion of how the stochastic conditioning affects the effective collaboration gain to Section 6.2 and C.5 In most cases of interest, // is close to 1, reducing &oyg to the raw heterogeneity level doyj- Thus, the following discussion of رلyة can be understood as applying to dabj well.
Proposition 2 previews the phenomenon of affinity-base basnit variance reduction. Compared to indepen- dent learning, Proposition 2 achieves a convergence rate accelcrated by a factor of max(n 8oug capturing speedup from both federated collaboration and agent similarity. When agents have sim- ilar objectives (dor) S n n) this factor recovers the linear speedup n- from FL (Proposition I); with abundant collaborating agents (n 2 km objective affinity dominates variance reduction. Importantly, since dekj E [0, 1]. AMPCLL's worst-case complexity is always upper bounded by that of independent learning, ((x21=)), ensuring collaboration never degrades performance. agents have markedly different objectives (ōot) t 1), collaboration benefits vanish and AffPCL falls back (rnssitioon]soowcaaa that AffPCL seamlessly interpolates
to independent learning, expected. between FL and independent learning, offering full adaptivity without imposing artificial restrictions.
To provide intuitions for affinity-basee variance reduction, discuss three interpretations of AIIPCL.
Bias correction. (x -f (x?)) ) in (5) corrects the bias in the aggregated update direction gf (x?) to achieve personalization. Specifically, one can verify that E, [7]] = E, [g (at)]. In collaborative learning, agents want to leverage the aggregated update direction for its lower variance, but also need to correct its bias towards the central solution rather than the personalized solution. We remark that this bias correction is fundamentally different from those used in the heterogeneous FL literature (Karimireddy et al. 2020b] Mangold et al. 2024), which correct for local drift away from the central direction. In other words, our novel bias correction term is personalized for each agent.
Control variates. Although (f:)) ( is personalized and thus cannot benefit from federated averaging, it enjoys affinity-based variance reduction via control variates (Defazio al. 2014 Rubinstein & Kroese 2016). Specifically. gf (==") serves as control variate that positively correlates with the local update direction gf (z:), and thus reduces the variance of the overall update. Then, a low- variance version (sample average) of the control variate, of (x-)). ) is added to correct the introduced
5

---

Published as a conference paper at ICLR 2026
bias. The variance reduction effect scales with the correlation in the control variate, which in turn scales with the affinity between the local and central systems. This control variate perspective offers clear explanation for affinity- -baee variance reduction in AffPCL, and motivates our design choice of the bias correction term 9; - (af), which correlates with of (z;) through the underlying sample si. unlike other potential candidates (e.g., A]r 69) that would be nearly independent of sj.
Central-local decomposition. To perceive how this variance reduction scales with affinity, we can view (5) as performing central and local learning in parallel. The central learning happens at the server side, seeking unified point z=5m that solves the central system (3). and the local learning happens at the client side, solving the local residual system Ar. 1looc b' 8 Then, x x² doocc gives the personalized solution to (2). Specifically. 9 (z²) = = by drives the central learning and (x g. (x²) (z) - =!)) - (b' - 5 (s{) drives the local learning. Intuitively. the local residual system is simpler to solve when agent's s objective is close to the central one, leading to affinity-based variance reduction.Ided Identifying this low-complexity local residual system is key to AffPLL's success. If the local learning problem were as complex as the original one (c.g., fine-tuning), such a central-local decomposition would offer only marginal speedup over independent learning.
4 INTRODUCING ADAPTIVITY: CENTRAL OBJECTIVE ESTIMATION
Knowing the central objective amounts to knowing other agents' objectives, which may not be realistic practice. This section removes this assumption by enabling agents to adaptively learn the central objective while learning their personalized solutions. A practical challenge is that when the state space S is large or infinite, 6 becomes high- or infiittedimeenionnal and learning it inevitably dominates the overall complexity. To match the dimension of other system parameters, we consider a linear parametrization of the objective function: (s) = (ss)0 for all E [n]. where & E Rdddd is a feature embedding function such that sym(R, (s)) > O, and 0 E Rd is the weight. This structure covers finite state spaces as a special case; see Appendix C.2 for more discussion.
Interestingly, central objective estimation (COE) is a special case of heterogeneous FL in Section 2 0 60, where 4 = (s)) and Ba = Σ. b², (s) (s). Therefore, agents can federatedly estimate the central objective using the same algorithm in Section
(6)
m Σ (s))
0.b (B)) 8 = 0g arge " (0.), where gt
(s)))s
n
n2
Without loss of generality, we use normalized features [1 I for all s E S. With linear
parametrization, redefine the objective bound G, == max (max; ||0: [|0 and heterogeneity level dayy maX | e (2G,) e [o, 1], which imply the original bound and Definition Then, COE directly enjoys the same guarantee as in Proposition]] |
with replaced by Amin (sym
We denote 80 == (s))0; as the estimated central objective at time t. Then, agents use (p())) in place of 60 (5t) in AffPCL (5). asynchronously with COE in (7). This scheme enjoys the same convergence
guarantee as Proposition 2 alas proven in Appendix
INTRODUCING ENVIRONMENT HETEROGENEITY: IMPORTANCE CORRECTION
5.1 CENTRAL LEARNING REVISITED
Section 3 discusses the central-local decomposition of AMPCLL. With homogeneous environments, central learning happens implicitly by considering the dynamics the averaged decision variable x! = 1 Σm= z/, which converges to the solution 2;; 1 "" 2f to (3). However, this is no longer true with heterogeneous environment distributions (1r * n²)). because
x - - E (("")~"")) # E A (HEL) (A")
That is, as x converges to the personalized solution x for all i E [n]. their average will not converge to S invalidating the implicit central learning through =.0

---

Published as a conference paper at ICLR 2026
Fortunately, manage to show that if agents explicitly maintain unified central decision variable If (3é x?)) and update it federatedly using (4). then x; still converges to x. with the same convergence rate in Proposition] [I even in the presence of environment heterogeneity. We refer to this explicit approach as central decision learning (CDL). The same argument applies to COE in Section 4 i.e., 7 converges the solution to (6) with the same rate under heterogeneous environment distributions. We defer the proof to Appendices E and F Intuitively, this works because sample from the mixture distribution 2e is equivalent to first sampling an index i uniformly and then sampling s from p². Therefore, a federated update direction that equally weights local sample information from all agents unbiased towards the central solution.
Beyond algorithmic implications, we remark that in Sections 2 and the averaged decision variable naturally corresponds to a virtual system with parameters J Ei=1'' 11' and b = Σi=
1 bi. The affinity among agents directly translates to the affinity between each agent and this "central agent" with index 0. Environment heterogeneity perplexes this concept: who the "central agent" now, and does it inherit the affinity among agents? Our central objective characterization in (6) helps answer the first question by defining (s()) == (s)0r, and then the "central system" (3) corresponds to a virtual system with environment distribution ot and objective be (# b9). Pinpointing this relocated central agent is crucial for deriving agent-specific affinity- variance reduction in Section 6..
The second question is more subtle, drastically different objective be from all actual agents, even when the latter have similar objectives. For instance, ill-conditioned system can amplify a small Sobj (Definition [} such that | b' ||o reaches its maximum possible value 2G, for some agent i. This divergence, which also applies to the relationship between the central and personalized decision variables, presents fundamental challenge introduced by environment heterogeneity in achieving affinity-based variance reduction. Fortunately, our analysis reveals that what is crucial is the affinity in feature space, e.g., terms like (IA =!Il and ||.4' 20 which are well controlled by the raw affinity among actual agents. Please refer Lemmas D.I and D.2 in Appendix for more discussion.
the central agent, unlike the "averaged agent", can have
5.2 PCL WITH IMPORTANCE CORRECTION
We now arrive at the most general setup (2), where agents have heterogeneous environment dis- tributions and objectives and seek personalized solutions. In addition to the challenges posed by environmental heterogeneity discussed in Section [5.]] a further obstacle emerges in the design of AffPCL: the bias correction mechanism Section 3 alooe no longer sufficient. To overcome this, propose integrating a novel importance correction to the central update direction before it gets sent to each agent, resulting in the following AfPCL update rule:
# = + orDi where
g = gf (x)) + dr (x - c+1 (x)),
(
where the bias correction term (z) -(())==8(s$} now uses the estimated central objective (s())) from COE and D. is the importance-correeted central update direction:
-
(((s)) ((s))x- (s))
(
(x)
12
AfIPCL (8) with asynchronous COE 7 and CDL gives the complete algorithm for solving 2 We provide the pseudocode and discuss implementation details in Appendix C.I
AfIPCL effectively handles environment heterogeneity by (i) correcting bias: E[g; 34 (x) (Lemma G.I). and (ii) reducing variance based on agents' environment affinity (Lemmas G.2 and G.3)
(x)
Definition 2 (Environment heterogeneity). The environment heterogeneity level is defined
daw max || pel || [0. (.j[(n]
The interpretations discussed in Section 3 still account for a portion of the variance reduction effect, especially w.r.t. objective affinity. For the newly introduced importance-corrected central update direction dt its variance has three key properties: (i) it decomposes into a federated term, an depeedet term similar to the variance of the bias correction term, and a term that characterizes the environment heterogeneity: d² x2(,, s14),, where x² is the chi-square divergence;
7

---

Published as a conference paper at ICLR 2026
(ii) the chi-square divergence is bounded by the total variation distance, which defines the environment heterogeneity: X²(n,, n²) & max{llp'Il lloo. 1:dews (ii)) the density ratio p has a natural upper bound: p' 1 ( دド= pe (s)) 1 (²(ss/ (s)) n. Combining the three observations gives an
upper bound of the additional variance from environment heterogeneity: x μ²) S o² This means that our method automatically adapts to the level of environment heterogeneity, enjoys affinity-based variance reduction, and never performs worse than independent learning, since dew 1.
These observations motivate the design of server-side importance correction. If this correction were performed on the client side, the additional variance term in (i) would lack the mitigating n- factor, and the density ratio μ (s)/ p0(p//(()) (s) would not be bounded by n as in (iii) which could result in potentially unbounded variance that degrades performance.
We are now ready to present the main result, which shows that AffPLL achieves affinity- variance reduction characterized by both environment and objective affinities, generalizing Section
Theorem 1. With a constant step size a Int Int/(At), AfPCL (8) with COE 77) and CDL (4) satisfies
E| x111 = Ö(x²t~1 max(n Saw Ho
Vi E [n].
where 8.m S m (, }. Juny minjsmin(11 (1. | are effective environment and objective heterogeneity.
6 DISCUSSION
6.1 DENSITY RATIO ESTIMATION
Section 5.2 requires that the density ratios (s() = [*(x)
(a(a) of environment distributions are known
to the central server. This is common assumption in supervised learning (Cortes et al. 2010
Ma et al. )023), controlled sampling (Rubinstein & Kroese 2016). and off-policy reinforcement
learning Precup et al., 2000; TTomass & Brunskill| 2016)) It satisfied, for example, when data
pre-collected or the covariate shift is induced by known mechanisms. When p is unknown, we can incorporate asynchronous density ratio estimation (DRE) into AffPCL. Similar COE in Section 4 DRE with linear parametrization (Sugiyama et al. 2012 Zhang et al. (2025) is also special variant of (2) (see Appendix C.4). However, unlike COE, which enjoys affinity
variance reduction
without importance correction, DRE seeks personalized solutions, which, according to our previous analysis, requires known density ratio for importance correction to achieve affinity- based variance reduction. This creates a chicken-and -andeeg-nnd-eg problem, settled by the following lower bound.
cermmation-thhoredee
Theorem 2. Let Pi be any estimator of the true density ratio p², given agents, t independent samples per agent, and no communication or computation constraint. There exists a system such that
inf En pb(() - p'))) 2 min ((96))-1,
Theorem 2 rules out the possibility of achieving variance reduction linear in the environment hetero- geneity level 6 without knowing the density ratio priori. This hardness result can be circumvented if additional structure presents, such as sparsity (environment distributions differ only in a few dimensions) or coupling (environment distributions are dependent). That is, the key difference from previous problems is that affinity in DRE should be measured by criteria other than total variation distance. Our analysis of AffPCL assumes access to DRE oracle capable of exploiting such structure to achieve affinity-based variance reduction, thereby proving Theorem i full generality and further showcasing the adaptivity of AffPCL. Appendix C.3 proves Theorem 2 and Appendix C.44 contains an extended discussion on DRE in our setting.
6.2 NOISE ALIGNMENT
We formally define the effective heterogeneity levels in Proposition 2 and Theorem as 8 min(1,vōaw) } and ب min (1, vo J. where // characterizes the system's 's "stochastic conditioning".
Definition (Stochastic condition number). V == max; (A') || where D(s) = AAss)T(ss),
v is trivially upper bounded by x (see Appendix C.5). and we refer to 1ー 2 -1 as the noise alignment constant. Note that the polar decomposition gives A(s) (ss) where the positive
8

---

Published as a conference paper at ICLR 2026
semidefinite matrix (s) defined above stretches the vector it acts on, and the orthogonal matrix U(s) rotates it. If (s) maintains a similar orientation for almost all s E S, A(s) is "well-aligned" and one can see that v' is large. Conversely, if (s) varies significantly, v tends to be small. The impact of affinity on variance reduction is thus modulated by this noise alignment (Lemma D.2)
We remark that while bears resemblance to the matrix condition number and is upper bounded by 66 the condition number и pertains solely to the deterministic parameters of the system, whereas captures the conditioning of the system's 's stochastic structure. Consequently, large does not necessarily imply a large v, and vice versa. Fortunately, in many cases of interest, the noise alignment constant vー is large. A particularly relevant example is a positive semidefinite A(s), a property often imposed on feature embedding matrices by design. In this case, V = 1. See Appendix C.5 for three more examples with large =1 and further discussion on noise alignment.
6.3 AGENT-SPECIHE AFFINITY HEASET VARIANCE REDUCTION
For ease of exposition, previous sections use the worst-ugent heterogeneity levels (Definitions T and 2) to characterize the worst-agent performance (Theorem [T. Intuitively, agents closer to the "center" should enjoy greater affinity- -aasdd variance reduction. In Appendix we analyze AffPCL in full generality and obtain an agent-specific convergence guarantee:
E| 11 = (((H))(() max B...)) Vi E [m]]
where n² o/\ X Amin (sim (A 82mm = min sm and &eo is more natural measure of agent i's closeness to the "center agent", defined as
8'm := max || - p ITv, 118 - - (l1l E [0,
Notably, 82m, is affected by both objective and environment heterogeneity, and admits a trivial bound 8'm minn(1,6e + days] (Lemma D.I). (9) confirms that AffPCL inherently offers agent-specifer affinity-based variance reduction, with agents closer to the center benefiting more from collaboration. An interesting consequence is that in the high heterogeneity regime, agent that is not close to any other actual agent (6am R 6ekg R 1) may still get a "free ride" by being close to the virtual central agent (5 < 1), thereby gaining significant speedup. Taking this a step further, an agent can collaborate with agents that are arbitrarily heterogeneous to it but still benefit from collaboration maximally and obtain linear speedup when 8 I n n-1 These insights are not captured by works that focus on linear speedup only in the low heterogeneity regime (Chayti et al. 2021 Even et al., (022))
7 NUMERICAL SIMULATIONS
Synthetic data. We first compare AIICLL, independent learning, federated averaging (McMahan al., 2017 FedAvg), fine-tuning (FedAvg followed by local independent learning), regularized
(T Dinh et al. 2020 Li et al. 2021 pFedMe and Ditto), and clustered (Ghosh et al.) 2020) FL methods, in synthetic system with 20 agents at different heterogeneity levels 8,m = day E {0, 0.05, 0.3, 0.8).
with results presented in Figure The average MSE = 1 Σi= MSE' is reported. For AffPCL, we also report the MSE of the agent closest to the center to highlight the agent-specific speedup effect.
Real-world data. We further evaluate AtIPCL on the real-world FEMNIST dataset. For clarity, we use only independent learning and FedAvg baselines. To introduce objective heterogeneity, consider a task where each user determines a handwritten character a digit and if it is a curved letter, with different users potentially having different preferences these two objectives. We train 10 users across four levels of objective heterogeneity and report the average test MSE in Figure
Reinforcement learning. We extend our method to the fundamental reinforcement learning algorithm SARSA, a temporal difference method that encompasses TD(O); see Appendix [.6 details. SARSA solves a non-linear policy optimization problem, showcasing the versatility of AffPCL beyond linear systems. We consider 10 agents with different reward functions (objectives) and transition kernels (environments). We incorporate in this experiment the asynchronous DRE module discussed in [[(s)]](aa)))
Section 6.1 Specifically, we have (i(s. a) : D (a)) where (s))) is the estimated state distribution of agent at time t via naive Monte Carlo, and 7/(a s) is the behavior policy of agent i at time l. The average MSE with respect to the optimal Q-function is reported in Figure

---

Published as a conference paper at ICLR 2026
Hornageneous
Low Heterageneity
Medium Heterogeneity
High Heterogeneity
Mean Squared Error
10
10F
10-4
10a4
10"2
30'
20
40
60
23
60
Samples
60
0
20
40
60 20 60
Independent Federated
PFedMe
Clustered
PCL
Ditto
Fine-tune
Agent-specific PCL
Figure 1: AffPCL matches FedAvg in the homogencous setting and independent learning in the high heterogeneity regime. Across all scenarios, AffPCL consistently achieves the lowest MSE, while other methods' relative performance varies with the heterogeneity level. In the high heterogeneity regime where all agents are dissimilar, the agent closest to the center still enjoys significant speedup.
Homogeneous
Low Heterogeneity
Medium Heterogeneity
High Heterogeneity
0.30
0.50
Test MSE
0.40
0.40
0.50
0.10
200
300
100
200
300
# Rounds
Indep
FedAvg
ArPCCL
Figure 2: AffPCL matches FedAvg in the homogeneous setting and consistently achieves the lowest test MSE across all heterogeneity levels. The relative performance of FedAvg and independent learning varies with the heterogeneity level.
Homogeneous
Heterogencity
Medium Heterogeneity
High Heterogencity
Mean Squared Error
104
10f
10-4
10:
10r4
1000 2000
3000
4000
1000
2000
3000
8000
1000 2000
3000
4000
1000 2000
3000
6000
# Samples
FedAvg
ARPCL
Figure 3: Consistent with other experiments, AtIPCL achieves the lowest MSE across all heterogeneity levels in the reinforcement learning setting. Incorporating asynchronous DRE does not hinder the performance of AffPCL, suggesting that density estimation is of relatively low complexity compared to policy optimization.
Please refer to Appendix B for the detailed setup and additional results. These simulations validate the superiority and practicality of AtPCL and our theory of adaptive affinity-basee variance reduction.
CONCLUSION AND FUTURE DIRECTIONS
AffPCL affirms that collaboration among arbitrarily heterogeneous agents can yield fully personalized solutions with adaptive affinity-baae speedup, opening new avenues for harmonizing personalization and collaboration in multi-agent learning. We advocate for future endeavors in the following topics: (i) personalized feature embeddings; (ii) trade-offs between collaboration benefit and communication sophistication such as cost. privacy, and security; (ii) lower bounds information exchange to achieve collaborative speedup: (iv) nonlinear systems, regret minimization, and stochastic optimiza- tion problems; and (v) other affinity structures such sparsity, correlation, and low-rankness.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by Jane Street, the MIT-IBM Watson Al Lab, the MIT-Amazon Science Hub, the MIT--ooggle Program for Computing Innovation, and MathWorks
REFERENCES
Sarabjot Singh Anand and Bamshad Mobasher. Intelligent techniques for web personalization. In IJCAI Workshop on Intelligent Techniques for Web Personalization, pp. 1-36. Springer, 2003.
Navid Azizan-Ruhi, Farshad Lahouti, Amir Salman Avestimehr, and Babak Hassibi. Distributed solu- tion of large-scale linear systems via accelerated projection-basee -astddd consensus. IEEE Transactions on Signal Processing, 67(14):3806-38) 2019.
Jalaj Bhandari, Daniel Russo, and Raghav Singal. A finite time analysis of temporal difference learning with linear function approximation. In Conference on Learning T'heory, pp. 1691-1692.
PMLR, 2018.
Avinandan Bose, Zhihan Xiong, Yuejie Chi, Simon Shaolei Du, Lin Xiao, and Maryam Fazel. LoRe: Personalizing LLMs via low-rank reward modeling. arXiv preprint arXiv:2504. 00.5504.44433 /4439, 2025.
Christopher Briggs, Zhong Fan, and Peter Andras. Federated learning with hierarchical clustering of local updates to improve training on non- con-IID data. In 2020 international joint conference on neural networks (IJCNN), pp. 1-9. IEEE, 2020.
Zheng Chai, Ahsan Ali, Syed Zawad, Stacey Truex, Ali Anwar, Nathalie Baracaldo, Yi Zhou, Heiko Ludwig, Feng Yan, and Yue Cheng. TiFL: A tier-based federated learning system. In Proceedings of the 29th international symposium on high-performunce parallel and distributed computing, pp. 125 136, 2020.
El Mahdi Chayti, Sai Praneeth Karimireddy, Sebastian U Stich, Nicolas Flammarion, and Martin Jaggi. Linear speedup personalized collaborative learning. arXiv preprint urXiv:2111.0500 2021.
Zhen Chen, Chen Yang, Mcilu Zhu, Zhe Peng, and Yixuan Yuan. Personalized retrogress federated learning toward imbalanced medical data. IEEE Transactions on Medical Imaging. 41 (12):3663-3674, 2022.
Gary Cheng, Karan Chadha, and John Duchi. Fine-tuning is fine in federated learning. arXiv preprint 10X1::2008.07313, 3, 2021.
Corinna Cortes, Yishay Mansour, and Mehryar Mohri. Learning bounds for importance weighting. Advances in neural information processing systems, 23, 2010.
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances in neural information processing systems, 27, 2014.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003. 13661, 2020.
Mathieu Even, Laurent Massoulié, and Kevin Scaman. Sample optimality and all-for-all strategies in personalized federated and collaborative learning. arXiv preprint arXiv:2201.13097 2022.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with
theoretical guarantees: A model-agnostic meta-learning approach. Advances in neural information processing systems, 33:3557-3568, 2020.
Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng-Zhong Xu. FedDC: Federated learning with non-IID data via local drift decoupling and correction. Proceedings of the IEFF/CVF conference computer vision and pattern recognition, pp. 10112-10121, 2022.

---

Published as a conference paper at ICLR 2026
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in neural information processing systems, 33:19586-19597, 2020.
Margalit R Glasgow, Honglin Yuan, and Tengyu Ma. Sharp bounds for federated averaging (local SGD) and continuous perspective. In International Conference on Artificial Intelligence and Statistics, pp. 9050-9090. PMLR, 2022.
Nathaniel Good, J Ben Schafer, Joseph A Konstan, Al Borchers, Badrul Sarwar, Jon Herlocker, John Riedl, et al. Combining collaborative filtering with personal agents for better recommendations. Aaai/iaai, 439(10.5555):31544 11555119-955,,, 1999.
Felix Grimberg, Mary-Anne Hartley, Sai Karimireddy, and Martin Jaggi. Optimal model averaging: Towards personalized collaborative learning. arXiv preprint arXiv:2 10. /2946, 2021.
Yongxin Guo, Xiaoying Tang, and Tao Lin. FedBR: Improving federated learning heterogeneous data via local learning bias reduction. In International conference on machine learning, pp- 12034-12054. PMLR, 2023.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. arXiv preprint wwx.iiv0002.0556.. 2020.
Chao Huang, Chen Lv, Peng Hang. and Yang Xing. Toward safe and personalized autonomous driving: Decissoonmmking and motion control with DPF and CDT techniques. IEEE/ASME Transactions Mechatromics, 26(2):61 1-620, 2021.
Devansh Jalota, Karthik Gopalakrishnan, Navid Azizan, Ramesh Johari, and Marco Pavone. Online learning for traffic routing under unknown preferences. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 3210-3229. PMLR, 2023.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008.036600 2020a.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sehastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pp. 5132-5143, PMLR, 2020b.
Mohamed Koutheair Khribi, Mohamed Jemni, and Olfa Nasraoui. Automatic recommendations for c-learning personalization based web usage mining techniques and information retrieval. In 2008 Eighth IEEE International Conference on Advanced Learning Technologies, pp. 241-245. IEEE, 2008
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International conference on machine learning, pp. 6357-6366 PMLR, 2021.
Xinyu Li, Ruiyang Zhou, Zachary C Lipton, and Liu Leqi. Personalized language modeling from personalized human feedback. arXiv preprint urXiv:2402.05133 2024.
Zeyang Li and Navid Azizan. Safe multi-agent reinforcement learning with convergence to general- ized nash equilibrium. arXiv preprint arXiv:2411. 15036, 2024.
Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas Allen, Randy P Auerbach. David Brent, Ruslan Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01522 2020.
Salakhutdinov, and Louis-Philippe
Cong Ma, Reese Pathak, and Martin Wainwright. Optimally tackling covariate shift in RKHS-based nonparametric regression. The Annals of Statistics, 51(2):38-761 2023.
12

---

Published as a conference paper at ICLR 2026
Paul Mangold, Sergey Samsonov, Safwan Labbi, Ilya Levin, Reda Alami, Alexey Naumov, and Eric Moulines. SCAFFLSA: Taming heterogeneity in federated linear stochastic approximation and TD learning. Advances in Neural Information Processing Systems, 37:13927-13981 2024.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprini 2xXiv:2002.10611 2020.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effee learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp.
PMLR, 2017.
Yurii Nesterov. Introductory Lectures Convex Optimization: A Basic Course, volume 87. Springer Science & Business Media, 2013.
Amin Heyrani Nobari, Kaveh Alim, Ali ArjomandBigdeli, Akash Srivastava, Faez Ahmed, and Navid Azizan, Activation-informeed merging of large language models. Advances in neural information processing systems, 2025.
Krishna Pillutla, Kshitiz Malik,
Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning, pp. 77716-17758 PMLR, 2022.
Doina Precup, Richard S. Sutton, and Satinder Singh. Eligibility traces for off-policy policy evaluation. In International Conference on Machine Learning, 2000.
Spencer M Richards, Navid Azizan, Jean-Jacques Slotine, and Marco Pavone. Control-oriented meta-learning. The International Journal of Robotics Research, 42(10))777-797, 2023.
Reuven Y. Rubinstein and Dirk P. Kroese. Simulation and the Monte Carlo Method. John Wiley & Sons, 2016.
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural networks and learning systems, 32(8):3710-3722 2020.
Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density Ratio Estimation in Machine Learning. Cambridge University Press, 2012.
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with Moreau envelopes. Advances in neural information processing systems, 33:21394-21405, 2020.
Sunbochen Tang, Haoyuan Sun, and Navid Azizan. Meta-learning for adaptive control with auto- mated mirror descent. In Proceedings of the 7th Annual Learning for Dynamics &amp; Control Conference, volume 283 of Proceedings of Machine Learning Research, pp. 1025-1037. PMLR, 2025.
Tao Tang, Zhuoyang Han, Zhen Cai, Shuo Yu, Xiaokang Zhou, Taiwo Oseni, and Sajal K Das. Personalized federated graph learning on non-llD electronic health records. IEEE Transactions on Neural Networks and Learning Systems, (5(9):118433. 11856, 2024.
Philip Thomas and Emma Brunskill. Data-efficient off-policy policy evaluation for reinforcement learning. International conference on machine learning, pp. 2139-21449 PMLR, 2016.
Boris Velasevic, Rohit Parasnis, Christopher G. Brinton, and Navid Azizan. On the effects of data heterogeneity on the convergence rates of distributed linear system solvers. In 2023 62nd IEEE Conference on Decision and Control (CDC), pp. 8394-8399, 2023.
Han Wang, Aritra Mitra, Hamed Hassani, George J. Pappas, and James Anderson. Federated TD learning with linear function approximation under environmental heterogeneity. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich. Zhen Dai, Brian Bullins, Brendan Memahan, Ohad Shamir, and Nathan Srebro. local SGD better than minibatch SGD? In International Conference Machine Learning, pp. 10334-10343. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2026
Stanistaw Wožniak, Bartlomiej Koptyra, Arkadiusz Janz, Przemystaw Kazienko, and Jan Kocoń. Personalized large language models. In 2024 IEEE International Conference on Data Mining Workshops (ICDMW), pp. 511-520. IEEE, 2024.
Guojun Xiong, Shufan Wang, Daniel Jiang, and Jian Li. On the linear speedup of personalized federated reinforcement learning with shared representations. In The Thirteenth International Conference on Learning Representations, 2025.
Linlin You, Mai Hao, Jian Sun, Yunpeng Wang. Chunming Rong, Chau Yuen, Paolo Santi, and Carlo Ratti. Toward personalized autonomous transportation system: Vision, challenges, and solutions. The Innovation, 5(6), 2024.
Chenyu Zhang. Han Wang, Aritra Mitra, and James Anderson. Finite-time analysis of on-policy heterogeneous federated reinforcement learning. In The Twelfth International Conference Learning Representations, 2024.
Chenyu Zhang, Xu Chen, and Xuan Di. Stochastic semi-gradient descent for learning mean field games with population-aware function approximation. In The Thirteenth International Conference on Learning Representations, 2025.
14

---

Published as a conference paper at ICLR 2026
Appendix
Table of Contents
Organization of Appendix
16
A Notation
16
B Additional numerical simulations
17
Further discussions
20
Implementation details
20
Lincar paramctrization objective and density ratio
................................
21
Proof of Theorem[2 2
21
C.4
Density ratio estimation
23
Noise alignment
24
C.6
Application to reinforcement learning
26
Preliminary Lemmas
E Analysis of Central Objective Estimation
31
F Analysis of Central Decision Learning
33
G Analysis of Personalized Collaborative Learning
36
Proof of Theorem
41
15

---

Published as a conference paper at ICLR 2026
ORGANIZATION OF APPENDIX
The appendix is organized as follows. Notation and symbols are summarized in Appendix A The detailed setup and additional results of numerical simulations are provided in Appendix B Appendix/[Isuppeeneee supplementss omitted discussions in the main text.
The remaining sections are dedicated to proving Theorem Tin full generality, incorporating asyn- chronous density ratio estimation and spent-specific step sizes, and subsuming all other results in the main text. We first collect several useful lemmas in Appendix Dito facilitate later analysis. The three components of AffPCL are then examined in sequence: central objective estimation (Appendix E).
central decision learning (Appendix F). and personalized local learning (Appendix
A NOTATION
We summarize the key notation and symbols used throughout this paper in Tables [ andd 2 reiterate that the superscript denotes the averaged quantity across all agents, i.c., f for any quantity ʃ. Due to symmetry, 8 usually satisfies the same property as s² for all E [n].
We
i= f
Therefore, in addition to the notation [n). also use (n²) to denote {0, 1......... ... The overline denotes the mean quantity under the corresponding environment distribution, i.e., f = f for
any operation f'. We remark that the aggregation of the mean values is not necessarily equal to the mean under the aggregated environment distribution:
HE f() n
8
_ B (
ENp
However, we have two special cases where the equality holds: (i) fr fā for all i.j [ni] or (ii) n² μe' for all i,j E [n]] The superscript denotes the explicitly maintained central quantity that aims to bridge the above discrepancy, e.g., the central objective If and central decision variable at. Generally, f * f for any quantity f. but the equality may hold in the two special cases above.
For the ease of presentation, we use the following shorthand notation throughout the analysis: As, represents the optimality gap - 2, at time step & for any decision variable 2; for a function f on S. fi represents its evaluation the observation si, and f = A E. f(s(;; EEj == Exj=r= and E, = Ext=ne ; Ex,-, represents the conditional expectation given the history filtration Ft-1 that contains all the randomness up to time step f
We use ,, to denote the Loewner order and R, S to denote the asymptotic order as D0.
Table 1: Notation.
Notation
Description
[n]
(1.2....,n), (0, I.....n)
A
decision variable optimality gap - 2n i-th agent's 's quantity, and its realization time step (if random variable) evaluation at observation st (if function) averaged quantity across agents 1 > fo
f fi
yo
fe
explicitly maintained central quantity expected quantity under agent i's environment distribution Eyg (s(s)
fo
aggregated expected quantity 1 E.=1E, f"
Jt
estimation of f at time step l bias correction from aggregated/central update direction to agent i importance-corrected update direction from central to agent
g0+ got
9
16

---

Published as a conference paper at ICLR 2026
Table 2: Symbols.
Symbol
Description
Symbol
Description
A
feature matrix
step size
objective
Young's s inequality parameter
C
constant
X²
chi-square divergence
d
system dimension
heterogeneity level
Estimation emror
(SS)
probability measure space
F
Filtration
density ratio weight
update direction
Y
reward discount factor
G
system parameter bound
condition number
j.
agent index
minimal eigenvalue
C
Lyapunov function
environment distribution
number of agents state, state space time step
stochastic condition number feature map
S
Φ, Φ b V
t
density ratio
weight
problem scale/ variance proxy objective weight
2,2
decision variable
B ADDITIONAL NUMERICAL SIMULATIONS
Synthetic setup. We run our simulations in a synthetic multi-agent linear system with n = 20 agents in a dimensional space. Agents possess distinct multivariate Gaussian distributions 4 = -m,, Ia) as their environments, and their personalized objectives are given by linear models (s)s) 4 (s) We construct the stochastic feature embedding matrices A(s) and ((s) to have multiplicative noise structure (Example 7) A(s) (la
ss" Abase and 4(s) (Ia 6 882" P where Abase and & base are randomly generated positive definite matrices for each problem instance with condition numbers of (1). We set σд = 1 and 0.5 control the level of stochastic noise alignment. The reference personalized solutions x² are calculated using Monte Carlo estimation with 5000 samples.
To generate heterogeneous environments, we set m, = Seny Ca where v; a random unit vector and CA 4 satisfies that [0, Ia) CC Ia) 0.9. This ensures that the environment heterogeneity level goes to 1 as 6am approaches (Definition (2. Similarly, heterogeneous objectives generated by setting 0 Obase + dokguu, where u, random unit vector and Ohase NN Ia)- This construction ensures that the objective heterogeneity level is of (p(ōo)) (Definition I.. For reference, we fix the first agent's 's environment as f = (0, Ia) and 0! = Obsse making it close to the "center" when the number of agents is large.
We compare our proposed AffPCL algorithm against the following baselines:
Independent learning, where each agent learns own solution using its local data without communication.
Federated averaging (FedAvg), where all agents collaboratively learn a unified solution by averaging their update directions.
3. Fine-tuning, where agents first run FedAvg for 30 steps learn common model, and then
fine-tune their personalized models independently for another 30 steps;
4. Regularized FL, pFedMe (T Dinh et al. 2020) and Ditto (Li et al. 2021) specifically, where agents collaboratively learn a global model, and then learn personalized models by solving regularized objective that penalizes the distance to the global model. The regularization parameter is set to 15 for both methods.
5. Clustered FL Ghosh et al. 2020), where agents are iteratively clustered based on the similarity of their local systems, and then collaboratively learn a model within each cluster. The number of clusters is set to 10.
The results are reported in Figure I All algorithms are run for = 60 steps with fixed learning rate of 0.01. All experiments repeated for 10 runs, and report the mean squared error
17

---

Published as a conference paper at ICLR 2026
= along with the 90% confidence region. To showcase the agent-specific affinity-base based variance reduction, we also report the error of the first agent MSF. = || - =.
averaged over all agents MSE =
Σx= |||
Comparison with baselines. We evaluate the performance of all algorithms under different hetero- geneity levels (daw, doti). Figure reports the homogeneous setting (0.0, 0.0), low heterogeneity (0.05, 0.05). medium heterogeneity (0.2, 0.2). and high heterogeneity (0.5, 0.5). Results of exhaus- tive sweeps over (6awv: 5dy)) are presented in Figure 4 where we report the improvement of AffPCL over independent learning and FedAvg, measured by the average MSE over the last 10 time steps. We remark that when agents' environment distributions vary greatly (6 2 0.9), all algorithms experi- ence high variance and thus the results may not be statistically significant. Figure [agure|4aji demonstrates that AffPCL consistently outperforms independent learning, with the affinity-based speedup increasing as the heterogeneity level decreases. Figure 4 shows that AffPCL matches FedAvg in the homogeneous setting, while FedAvg fails to provide any personalization in the presence of heterogeneity.
160
100
0.9
6d
-00
457
RR
1C0
1
-11
4
RZ
18D
0.9
71
-11
4
FRI
1D
GO
12
25
0.8
28
32
26
24
16
32
16
28
15
0.8
41
53
28
53
50
33
53
25
25
25
80
80
0.7
33
46
31
25
32
40
30
16
18
15
0.7
55
82
78
53
55
26
26
97
0.6
43
42
50
44
48
31
29
20
22
40
In 1963
0.6
64
57
52
94
96
52
97
97
97
inorovement 1955
60
0.5
59
53
59
54
49
52
20
30
36
31
0.5
89
90
96
96
57
56
96
98
0.4
0.4
57
97
50
0.3
0.3
0.2
0.2
90
52
0.1
0.5
0.6
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
6g
(a) Over independent learning.
(b) Over FedAvg.
Figure 4: Improvement of AffPCL.
Federated vs. affinity-based speedup. Our theory identifies two factors in the variance reduction of AffPCL: federated speedup - and heterogeneity level 5. We conduct exhaustive sweeps over the number of agents E [2, 50] and heterogeneity level 6 = daw dayg [0.02, 0.5). Iso-performance contours of AtfPCL are plotted in Figure where each curve represents the combinations of (n (,=1,)) that yield the same average MSE over the last 10 steps. As expected, the contours form Pareto-type curves, confirming that max (n 11 .5) characterizes the trade-off between collaboration and affinity in AffPCL.
0.4
HIE
3e-
011
1e3
0.1
02
0.3
0.4
Figure 5: Iso-performance contours of AffPCL.
18

---

Published as a conference paper at ICLR 2026
Agent-specific performance. Another highlight of our theory is the agent-specific affinity- based vari- ance reduction effect, where agents closer to the center benefit more from collaboration (Section 6.3). We examine this phenomenon in high heterogeneity setting (ó dakg) (0.7, 0.7) and report the performance of independent learning and AffPCL for a generic agent and for the agent closest to the center in Figure [ Looking at the performance of generic agent, AffPCL performs similarly to independent learning, as the collaboration benefit gets diminished by the high heterogeneity. However, the agent closest to the center still gets a "free ride" and achieves significant speedup, compared to learning on its own, through collaborating with other agents. We remark that in the high heterogeneity regime, the agent closest the center may not be close to any other agents, yet the collaboration benefit remains.
Generic agent
Agent close to center
10c
10
MSE
10
10
10-2
10-2
20
40
60
20
40
60
Samples
Independent
PCL
Figure 6: Agent-specific performance.
Real-world setup. For the real-world FEMNIST dataset, we first pre-train a numerical network using some training data as the feature extractor 6 shared across all users. The dimensionality of the extracted feature is d 84. To introduce varying levels of objective heterogeneity, we consider the following label for input character image 2:
y(z; A) = A [character z is digit) (1 - A) 1 [character : is a curved letter). Then, the least squares linear regression problem corresponds to the following linear system:
E,i [o(=)p(2)²] x' E,,: [o(=))(2;;
,,,
where n² is the data distribution of user In implementation of AffPCL for this task, omit the importance correction. We test four levels of objective heterogeneity by setting the range of as [0.5 - dohi /, 0.5 + 6ak/2] with dobj {0, 0.2, 0.6, 1}. 10 users have evenly distributed X in the specified range. All three algorithms share the same hyperparameters: learning rate
0.002. batch size is 32, and the number of samples between two communication rounds is 70. The test MSE averaged over all users is reported in Figure The results are consistent with those in the synthetic experiments, validating the practicality of AffPCL.
Reinforcement learning setup. We follow the derivation in Appendix C.6land setup in Zhang et al. (2024) to implement AfIPCL, independent learning, and FedAvg versions of SARSA. We consider 10 agents, 10 states, 5 actions, and feature dimension of = 20 for the state-action space. The other RL hyperparameters are set as follows: reward discount factor 7 = 0.1, behavior policy is 0.1. We perturb nominal reward function and
softmax with temperature 10, and step size transition kernel to generate heterogeneous objectives and environments, with relative perturbation levels dakg, 8em E {0.0, 0.2, 0.5, 1.0). The reference optimal Q-function is calculated using offline value iteration with the true model. For the asynchronous DRE module, we use the following density ratio estimator:
(f(s)=f(a (s. (
where (s) is estimated using the naive Monte Carlo. i.e., average state visitation frequency up to time stept for agent i, and 7/ s) is the behavior policy of agent corresponding to its current Q- function estimate. The average MSE with respect the reference optimal Q-functions is reported in
19

---

Published as a conference paper at ICLR 2026
Figure 3
The results consistent with other experiments and show that incorporating asynchronous DRE does not hinder the performance of AffPCL.
C FURTHER DISCUSSIONS
C.I IMPLEMENTATION DETAILS
We present in Algorithm T the pseudocode of AffPLL with asynchronous COE and CDL.
Algorithm 1: Personalized collaborative learning (AffPCL)
initialize: =y. 0g ar for [n]-
for / 0, 1,... do
foreach agent [ra] in parallel do
sample of evaluate residuals gi(x))), (ert))..(((( (x(), gi (09), .. (af)
send (s), (f si (0f), (xf)) to the server at central server:
aggregate central residuals Ir (=f) 0.b (01) " for [n] send central residuals back to agents
0.0
foreach agent E [n]] in parallel do
x = x - atht (xf)
10
01+1 Og arg: 0.6 (09)
11
#[+1 = st (x(()))))) + ozi (xG) 9i c+i (x)))
12
provide several remarks on central decision learning (CDL). First, if central server has memory, the central decision variable zff and central objective parameter 0f can also be maintained and updated (Lines 10 IT)
at the server side.
Second, when agents share the same environment distribution, the central decision variable can be replaced by the average decision variable zff, since the solution to the central system (3) coincides with the average of personalized solutions, i.e., xº_ x²Q. In this way, CDL happens implicitly without executing (4). reducing computation complexity. However, this implementation requires an additional communication round to compute x0 after each local update, increasing communication complexity.
Third, when agents have different environment distributions, note that the central update direction c (af) in (8) now involves gf (x;) (which uses the estimated central objective by (s))), instead of (2f)) (which uses personalized objective b' (s{)) as in (4). This modification is necessary for the importance correction to work. That said, it should be unsurprising that if we also use of** to
compute the central update direction in central learning, i.e., using oe (z 1 E Dr (xxf) in (4). the convergence guarantee still holds. For completeness, we also prove convergence for both implementations in Appendix Fi
Remark 1 (Communication and computation complexity). To provide clearer picture of scalability in our stylized setup, we compare the communication and computation complexity of AffPCL with federated averaging (FedAvg), both with immediate communication after each local update. In FedAvg, each round involves the communication of all local residuals and one averaged residual, resulting in a communication complexity of comm comm(FedAvg)) system dimension. Updating the local decision variable results in computation complexity of comp(FedAvg) (d) per agent per round. In AffPCL with COE, CDL, and DRE, as detailed in Algorithm [] each round has communication complexity comm(AIPCL) < 4 comm (FedAwg)). Suppose we calculate the central update direction using of as discussed above, and suppose the density ratio is known priori, then the communication complexity reduces to comm/ AffPCL) < @(5nd). Similarly, the computation complexity per agent per round in AIIPCL
@(2nd) per round, where is the
2.5 comm ( (FedAvg) (Algorithm [] is comp(AfPCC AfIPCL) = 6 comp(FedAvg) = ((6d), where the extra factor comes from CDL, COE, importance correction, and composing the personalized update direction.
20

---

Published as a conference paper at ICLR 2026
The above analysis shows that, although sharing the same asymptotic communication and compu- tation complexity as FedAvg, AffPCL incurs a larger constant factor (specificallyy, up to 4 times for communication and 6 times for computation per iteration) due the additional components and personalization. This inherent trade-off calls for future research into developing more sophisticated communication schemes reduce the communication and computation overhead of AffPCL, and deriving lower bounds on the information exchange to achieve collaborative benefits in personalized learning. Additionally, we remark that numerous techniques from federated and decentralized learn- ing that reduce the number of communication rounds, such as multiple local updates, compression, and partial participation, can be readily incorporated into AffPCL. We consider the stylized setting of immediate communication to focus on the main ideas.
C.2 LINEAR PARAMETRIZATION OF OBJECTIVE AND DENSITY RATIO
We first note that linear parametrization covers finite state spaces as a special case. For DRE, (s)) = eg is the one-hot encoding of state and then n/. (s1), p' (s)ss ))" simply records the density ratio for all states. For COE, we transform the original objective as b'(s) < eg @ ((s) Rd|s|, where @ denotes the Kronecker product. Then, (s)) = ege @ la E Rd55s and 03 similarly records the objective vector for all states.
Linear parametrization also widely used in supervised learning (parametric regression) and re- inforcement learning (linear value function approximation). COE performs parametric estimation with linear function class: 0. argminggR= |lb - b'11, where (s) o(s)0, and omit the discussion of approximation error when the model is misspecified, i.c., b'(s) & { A e Rd). See (Sugyyama et al. (2012) for more details on DRE with linear parametrization.
When applied to reinforcement learning, our linear parametrization subsumes linear Markov reward processes Bhandari et al. 2018), where by (s) = (p(s)r" (s) = (s)), (s)T0', with r' as agent i's reward function and Op, Φr feature maps. See Appendix ( for more details on this application.
We choose linear parametrization for its simplicity, allowing us to focus on the main ideas. Our method readily extends to other (nrammraee models, provided the function class has complexity polynomial in d rather than in (S].
C.3 PROOF OF THEOREM 2
We restate and establish the lower bound of DRE MSE. Our proof generally follows the standard Le Cam's method with two special treatments: (i) we show that DRE is lower bounded by a special density estimation problem, (ii) we show that collaborating with other agents does not help in estimating one agent's own density ratio.
Theorem 2
Let (s(s) be the estimate of true density ratio p', given by any algorithm with agents and independent samples per agent, with no communication or computation constraint. There exists problem instance such that
inf pt b - p=) min (9961) 82 8_
Proof. In this proof, we omit the agent index and thus u/n². We build the hard problem instance step by step. We first consider a finite state space S = [d] and thus and Ot are vectors Be (this is equivalent to a linear function approximation with feature map (s) = e,, the s-th standard basis vector). Then, the minimax risk w.r.t. the MSE loss is defined as
R = inf sup Pe OSP5n [p B (e((s) p(s)
sup
where P.. a random vector, is the estimate learned from the sample drawn from (μ 10)01, We use the convention that if any constraint set empty, the risk is zero. We note that although samples across time steps are i.i.d., the samples from and nd at the same time step are correlated as μ = 1 Ey= p².. However, this correlation scales as On-)) and diminishes as n increases. We then apply several reductions. By set equivalence,
R" inf sup sup A p=pp SSSpS
E e Ep) p(s))
21

---

Published as a conference paper at ICLR 2026
We now fix 2t d ((S) and define convex set M 10 10 ppl E A(S))
Then,
R" 2 inf sup B(ux)o)e:B p(s))
Pe pmmpp,,EEM .EE
=d inf sup Bluxno0o:1lrr - pll PEM
For sufficiently small to be determined, we can choose pl, E M such that || - -22l2 2 2e. Following Le Cam's method, we construct 40= argmin 1 Pm 2|2. Then,
2
Rd 2d inf Pt 2
P (9t # m)
m==
2
2d 2 inf
Pm (f:t # m)
m=1
2 (1 - P1 - Pallrv),
(10)
2d
where Pm (mm p2)st, By Pinsker's inequality and properties of KL divergence.
P] - Palliv SDKI
(11 x 1ollue n2))
=t(D(((Hill)e) DKL (milleg (())) kn)).
(11)
H.
H2
where ph is conditional distribution of 80 given s ~ pm and 50) ~ pm x nb. Recall that 2 is the aggregated distribution of agents' distributions.
We proceed to construct Pl, P2 and bound the KL divergence terms. Suppose d is even. Let
(m()) = d + d 3/2 dd-))/((-1)-mm E [d,m E [2].
Let e < Vd. One can verify that pm E (S)) and
3/2 |l Pall2 d | Hallz = d 2d 3/2 c. Vd = 2c.
Further, let < vd/2. Then, for the first KI. divergence term, we bound it using the chi-square divergence:
(Hr(s) (2(s))
4d 32 d-1/2
Hi S 5x2(111122)
8d 1
(2(6)
For the second KL divergence term, we first have the decomposition of μrº 1,2 Um + n=1 H/m for m E [2]. where f',, is the aggregated distribution of all agents except agent i and is independent of Umm Then, the conditional distribution given s ~ pm is simply 1 8 + n-1
n pmm- for m E [2). n 1, then p = 12 and H2 = 0. Thus, we consider n 2 By the convexity of KL divergence and Jensen's inequality,
H₂ DKLL (f. 11,, (m) n DKL. (14111 [m)) n
DKL(HGllb))
n
Again, bounding it by chi-square divergence gives
-s) (b(())
H₂ < n
(s)
1=1
Notice that
0
72
1/
+
n
1/
(μi - 16)2 = 111- 111-11 142
3d-322
(n 1)2
HHere ona the element-wise product we treat them as functions S.
22

---

Published as a conference paper at ICLR 2026
Together with (c2(s) d=1 n 3/2 (c(-1)* d d have
8d 22
H2 x
n(n - 1) S 4d 12
Plugging H1 and H2 back into (TT and combining (TO) gives
6² 2d
61
R 2
11
(12)
We are left to determine There are two cases. The first is that dew S 4v6t' ".v0.i.. ie.. μ is close to い or we do not have many samples. In this case, one would constrain the estimator to get a smaller error. Note that
-
Mm Ir = d. d 3/2,, m E [2].
2
Thus, pushing jir and ugg to the boundary of the ball {w: : || 10 S daw gives 20 Vd. Plugging this into (12 gives
Rd 2 28 28 (1 20 v 6t) 2 em
The second case is that deny
1/ where we need make ln and μ22 closer make their
discrimination harder. In this case, we set e /V or Then ll p
dew So μr and av
μ22 are feasible. Plugging into (12) gives
R' 1
96t
Combining the two cases gives
R 2 min p (96t) -1.
C.4 DENSITY RATIO ESTIMATION
This subsection first shows that DRE with linear parametrization a special variant of and then discusses several environment affinity structures that help circumvent the lower bound in Theorem and enable affinity- based variance reduction in DRE.
A linear parametrization of density ratio (Sugiyama et al. 2012) takes the form p'(s) = (ss) n for all E [n], where g)(s) IR. is a measure basis and 1/ E Rd is the true weight. Let (+) = (s)) Then,
E Vs))t
= (s)) (
(s))
Therefore, a simple stochastic fixed point iteration for (2) described in the paper (see also Example 2 finds v/ with an MSE of O(t-1), but the affinity-based variance reduction is unattainable because of Theorem
Alternatively, notice that ps(s directly measures the affinity between j² and 1a and is the quantity through which p enters the analysis. Thus, we can directly apply linear parametrization to p' (s) - 1 (s)
TT n/. Then at time step t, p = 1 + (bt)) nt. The DRE problem becomes
E (
(()) (p' (p)(((
(s)
(55) (n' (s)) ds Ep--po o (s).
This problem formulation is easier to work with because n/ when js¹ nº: in such cases, regularization can be applied it is known a priori that agents' environments are similar.
Theoretically, this regularization will not work if our prior knowledge of environment similarity is measured in total variation distance, i.c., 8encs due to Theorem] 2 Nonetheless, several other affinity structures can help.
23

---

Published as a conference paper at ICLR 2026
Example (Sparsity). ny. encodes the difference between p' and n (recall that in the tabular case, (s) = ps() - - (s) (s); see Appendix C.2). If μ and μt differ only in a use fo-constrained
few dimensions, i.e., vi. is 6em d-sarse such that |lo I Seny d, then or f1-regularized least squares to estimate v/. which can be calculated efficiently online and has standard MSE of Ōxx2t=1_ d).
Example 2 (Coupling). Suppose DRE uses coupled samples from j and n such that P(s) = = ll 1rv and the equality is attained when je' and μ
1 - bun- Note that P(s s
optimally coupled. Then, simple fixed-point iteration
Mht = vi algi -0 (m²) = ni of (v())) (W(ss) - vb()))))
with properly chosen step size af = Ō(t has MSE of Ō -Baw)
Proof. We only need to show that the update is monotone and its variance at the fixed point enjoys affinity-based variance reduction; then the result follows from standard stochastic approximation analysis (see e.g., Appendix Eb. First, we have
Egi 1. (m)) = E (s))ni - E (s)) = "
m²) --
The montonicity follows from
Anii T An))) 2 Amia
For the variance, without loss of generality, we assume normalized feature [|v(())] 1. Then,
Ellgg 6-P (v/)1 =E =EI (1(s)) (
=E| (p)) -) (w(s)) ((s)))
s (10( + 2E]()(()) - (+())1?
< 2x (x', +2 4P +
<2 22 μ Irv + 8(1 - P(s =())
(65m))
These examples indicate that DRE requires a stricter affinity measure to achieve affinity-based variance reduction. To enable maximum generality, make the following assumption.
Assumption 1 (DRE oracle). We assume access to a DRE oracle that returns an estimate weight vi or density ratio o such that (p{(s) p² (s)) O(1) throughout the learning process and
E () - (s) =()())))))) max(n -1 1, %.8mm }).
(13)
where Rº captures the conditioning of the DRE problem.
Assumption ennsures that DRE does not become a bottleneck for achieving affinity- based speedup. Our analysis incorporates asynchronous DRE through Assumption [ see Appendix Gi
C.5 NOISE ALIGNMENT
We first establish the trivial bound of the stochastic condition number (noise alignment constant) defined in Definition3 3, In this subsection, omit the agent index for simplicity and generality. For a general stochastic matrix A(x), recall its stochastic condition number is
= | '11
where Ā EA(s), D = Es))), and D(s) = / AA())* A(s). For the upper bound, we have
V S ||
The "numerator" satisfies
||||| [+)( E (()) Ex (D(s)) Ea (A(m)(((s))) GA
24

---

Published as a conference paper at ICLR 2026
where we use the polar decomposition A(s) = (ss) where (s) is orthogonal matrix, and thus A(s) and (ss) share the same singular values. The "denominator" satisfies ||A-'| o min (A), and we have
minu (A) = min
|Ax|l = min
||
min x" Ax min x² sym(A)æ Amin (sym (A)). ||
|
||=||1
(14)
Note that À 2 Amin (sym(A)) and GA S a. Thus,
v V
To illustrate the idea that wー measures the alignment of noise in A(s), we describe one example where v is small.
Example (Misaligned noise). Suppose S = 10, 2n - e] C R and A(o) = U(s) E R2x2, where U()) is rotation matrix with angle Then, D(s) I, D and
A
25
COS s sin
cos E - sin ds sin cos
cos - sin 1)
2 sin
cos sin 2
sin cos
whose smallest singular value is sin 2 when O is small. Thus,
and wー 0 as c c-> 0. This is an example where the orientation of A(s) is uniformly random, and thus the noise is completely misaligned.
We then give several examples where the noise is well-aligned and the stochastic condition number equals or is close 1.
Example 4 (Constant orientation). Suppose A(s) U ((s) for all S, where U is constant orthogonal matrix and ((s) ł 0. Then, A = UD and v = 1.
Example (Positive semi-definite matrix). Suppose A(s) ł for all S. Then, D(s) A(s) and y = 1.
Example (Low rank feature embedding), Suppose the feature embedding matrix A(s) has low- rank structure: A(s) (o(s) - r (s), where ((s), (s) RR are two normalized feature maps such that [[((s] = [+)[l = 1 and dos) d (bss) for all E S. This is the case of temporal difference learning with linear function approximation [Bhandari et al. 2018), where 7 is the reward discount factor. Then, 1+7
Proof. We have
2 = (6())T(()) 2700() + +70))))))))))))) (())
- (00)))()(((()
which implies
D(s) (1((s) mw6()) ((s))(s)² = (1 ()))))(()
|[0())]
On the other hand,
A(s) 2 (()))(8) (0()))())) (())(())7) = (1-2))(()6(())
Thus,
Az - (((1=))()()0000))) 1-2 D,
1+7
which gives
1+ + DA-1
1
Thus, v S
14 T-
Remark 2. While strict normalization [|o()) [s)]| is required to show < 1+y
1-y* That is,
Lemma D.2 the only place where v is directly used, also holds with [|o(s)] [(()) < define for interpretability and ease of calculation, but can be relaxed in certain cases.
25

---

Published as a conference paper at ICLR 2026
Example 7 (Multiplicative noise). Suppose the noise is multiplicative: A(s) (I + U()) where (() is zero-mean and [|U(s)| Σ for all s E S. Then, v S +
Proof. We have
D(s)² D03--- A''
= A
Finally, we refer readers to Lemma which illustrates how noise alignment affects the translation of the affinity into an effective affinity for variance reduction. Intuitively, when Á is ill-conditioned, a small perturbation in the objective can lead to large perturbation the solution. Then, A(s) does not align well with A, the large perturbation in the solution may be further amplified by the large eigenvalues of (s)). leading to a large variance in the update direction. On the other hand, when (()) aligns well with A, the perturbation in the solution is only stretched by the small eigenvalues in A(s). maintaining similar magnitude to the objective perturbation and preventing noise amplification in the ill-conditioned subspace.
C.6 APPLICATION TO REINFORCEMENT LEARNING
This subsection gives a concrete application of our method to the policy evaluation problem in reinforcement learning (RL) resulting in personalized collaborative temporal difference (TD) learning.
Heterogeneous federated RI. has garnered traction recently (Zhang et al.] 2024 Wang et al. 2024 Xiong al. 2025 Li & Azizan 2024) due to its practicality by accommodating heterogeneity in multi-agent decision decision-making. However, existing works either fail to personalize and hence only work well in low heterogeneity regimes (Wang et al. 2024 Zhang et al. 2024), or deliver slower convergence rates (Xiong et al. (205). Our framework encompasses the setting of heterogeneous federated RL and method provides the first personalized collaborative reinforcement learning algorithm that accommodates arbitrary heterogeneous agents while achieving affinity-based variance reduction.
Consider n agents with distinct Markov reward processes (O, Pi, Ri, 7), where O is the state space, p is the transition kernel induced by agent i's behavior policy, Rs O O + IR is the reward function, and E [0. 1) is the discount factor. Following (Bhandari et al. 2018). we write (o) E[R (oh, 1)) 1) 0j = o). Agents want to evaluate their behavior policies by calculating their infinite horizon value functions v'(s) = E E > (o),) of = 0]. where oi ~ ( of)) With a linear function approximation (o) so Ø(o)² rf for some ri E Idd, the expected projected Bellman equation can be cast into as
E (d( + E (s((s))
(15)
A(s.s")
where E' ( The stochastic residual of (1)) is the TD error, and the corre- sponding fixed point iteration gives the TD(0) algorithm. Specifically within our framework, each observation tuple is s (o),. oh A({)) ((s)) = ((o (o(o,, yr y0(of +1")., (s()) o(o),, ) (f, and the environment distribution is (0, o') = (0) x P'(d' lo), where nं is the stationary distribution of the agent i's transition kernel. Then (x() represents the TD error and AffPCL (8) gives personalized collaborative TD(0).
With a normalized feature map [|*(o)]| < 1 and constants G2 >
(n)) || | G3 maxieja] |/ we have GA < + < 2 and < 2 max (1 + G), Go < G + Go.
As shown in Bhandari et al. 20188 Lemma 3), X > (1 Amin (E (d(s)(ss))) As shown in
Herr assume offline RI. setting where we have i.i.d. samples from pre-collected dataset consisting of observation tuples (of, t P- (o)), R (oh, o)))))
26

---

Published as a conference paper at ICLR 2026
Example 6 the stochastic condition number in this example is by Theorem I the sample complexity of personalized collaborative TD(0) reads
1+2 - 2(1 (11-))_1 Therefore,
1-7
o (G2 + G.)² (1 7 (1-7)8 max(n Semv, 6ass)
where w¹ = Amin (E (6
and den day represent kernel and reward heterogeneity levels, respectively. This complexity matches the best known result for homogeneous federated TD learning (Wang al. 2024). while offering new insights in high heterogeneity regimes.
D PRELIMINARY LEMMAS
Lemma D.1 (Affinity). Given the universal scores daw = max;,; I| - je ||rv and dekj = max,, |00 (2 along with the agent-specifie scores 8am = I1 - pe l|rv and 8āg = |00 (2G), we establish bounds on various parameter differences in terms of these scores.
(a) [bb(ss (s)]l < 2Go8ogg for any i,j E [n²].
(b) [ 8 (s)] - 2Gro ai for any E [m].
(c) l B'1 S 2Gm + dekg), for any i,j [n-).
(d) || 1 S 2G, min (1, 82 + Sobj 8am b. for any i [n].
We thus define 6'om = max { (5a 1 -1)1 & min (0, 6m + Seng, Sem + +
(e) | (8) b*))]] < 20 where E (0, for anr i E [m].
( (Naive) 10 0.1 < 22 22- ( + ddgg), for any E [n].
(8) | 1 S A for any i E [r="].
(h) | 1011 s A. for any i [m].
(i) (A)) - >1]]]1 S 20 where j E {0, i), for any i E [n].
(j) (Naive) | - =xil s 2x 0 + far). any i [n].
Proof. For Item (a) by the linear parametrization the objective,
[1b(()) - (s)]l = ()) 0iل1| < 1 0.1 < 2Go
(16)
where we use the fact that [b)) < 1 and the definition of daij in Section 4 the same argument with agent-specific score d'dej used.
Item (b) follows from
For any function f such that [f][l < G for all E S, and for all i E [n], by Definition 2
((s f()]l - 11/ -s)) (²
S2G || - 1LllTV
2G jon j [n]
$
(17)
2G 2G
This bound first gives Item (g) and Item (h) by letting f(x) = A(s) and G, = GA- Then, combining
16 and 17 with f(s) (s) and G = G, gives Item (c)
b')
E
E R (+)] 2G 2Grd.mm 2Go(faaj + dem))
27

---

Published as a conference paper at ICLR 2026
Specifically for the difference between the personalized and central expected objectives, have
118 1
(s)
1
(()
j=1
(s) 1 Σ
(E Es)(s)++ +
j=1
<1 - 60 (s()]l + Σ
I(E (+))
$2Gbo + 2Gr
Similarly, we have
(s)(ss 1
| 59
n Σ
BBbl(ss)
1 Σ - +
<| - E R (b'(
<2Gbbobj + 2G
The above two bounds give Item (d)
We then look at the naive bounds Items (f) and (i) on the difference between optimal solutions. For
any j E [n]] we have
A' - x!) + (A' A (5' B2 = 0.
which gives
- S II(AA) 112 (IAA* 11 + 11b R 12112)
Combining the previous bounds on system parameter differences gives
||~:; = (2' G2 2Gr( + 8mm)) T (At) 20 (bab) + daw))
The above bound also holds for the difference between the personalized solution and the central solution satisfying A b0. Specifically, the same argument gives
|| -xil22 S Gmiia (A") 20 (5org + 6aw)).
Let À minjeja] min (umin (sym( Amiu (sym (ū')) (14) gives Item and similar argu- ment gives Item (f) Notably, the upper bound of the optimal solution difference scales with X-1 which can be large when A' is ill-conditioned. This indicates that the affinity in objectives or environments do not translate well to the affinity in optimal solutions.
Fortunately, the bound is tamer when the optimal solutions are left-applied by A::
|| - 12) =11A's - Ao + (Ā'
A²)x = 1| - (Ā' - A²)x
22G cm + 5 G2 < 20 mm
Left-applying A gives the same bound, thus giving Item Item (e) can be derived similarly. Items (e) and (i) are saying that the affinity is well-preserved in the feature space, i.e., the image of the feature embedding matrix. This is also key to our analysis: we will never directly bound the difference between optimal solutions, but always inspect them in the feature space.
Lemma D.2 (Effective affinity). Denote 8imm = min { 18 ). Then,
Ey A(s x2 < 28 & vi E [n].
E, (s)) 02)11² < 20 8im.
28

---

Published as a conference paper at ICLR 2026
Proof. We only prove the first inequality; the second be proved similarly. We have
E'IlA(s)(a+ - =0)l² - " (A)) A()) - x
-(x; " P -x
<|10; > = |E (+)) (s =:)1l
2G GG / [ x1I
2G D (A') A's -
2G AG A xiIl
SGG 20
where the last inequality follows from Item in Lemma have
D.I On the other hand, by the trivial bound,
(')lA(s)(d) - S (2GAG#)².
Combining the two bounds gives
(=11A(s)<1==()) 20² min {1, vd } S 20 20 cen
Lemma D.3. For any two distributions pl over S, we have
x2,,1 S max (llp///loo: 13 l1 p a/ll
Proof. Let Sq = {s E S: (ss) 2 (l()) We have
(u, /
((s)
n² ds ds
(((s))
su(s) n((s))
1 (s)
So
((s)
(u(+) 11)))) ds /
11_ mu(s)
(1'(s)
m(s))
ds
(s)
(s)
So
20
20,51
<(max (llm//' 1o,, 1] - 1) 1 || + 11 11'l//
max
{ [1ll// Io,, 1} ||I μ'llTv.
Lemma D.4 (Uni-timescale Lyapunov analysis for asynchronous learning). Consider asynchronous learning of multiple decision variables =ti &s E [K], which satisfies the following one-step contraction:
k-1
BP < 31-3
(a + of Σ Ck CK*
== ==1...,, K.
That the convergence of z also depends on the other decision variables zK² W' < k. of is the step size for z;; we set them using a unified effective step size Oy:
at =a <1
k=
k=
K
w
1
1,
w
Σ
w M (2
k k 1
k==+11
Consider the following overall Lyapunov function:
L.
E A
Then, we have
K
L1t1 S (1 - at + a Σ
w (2
k=
29

---

Published as a conference paper at ICLR 2026
Proof. By definition, have
K
L1+1
w * E
K
w
(1- (1- 11-32at")B1A>3 + (a + +
k B|| - o Ip)
w0((1 )aa +
Σ
wn ai CK'LK B|| + wn (at)pC*)
(18)
(w (1 ( + w a. ( |1|A2111 w +
(19)
((1 a B (at)
k
=(1 - 04 ) + 0
Σ
w C** (X
where 18) follows from rearranging the summation and grouping the coefficients of E A [| and 199 follows from the definition of w and a.
Lemma D.5 (Constant and diminishing step size). Suppose we have the following one-step contrac- tion:
L (1- + a
Then, with constant step size Int/t, we have
L. = 0 Clil
t
With a linearly diminishing step size a, 4/((T + to 1)), O,...,t, the following convex
combination
L. Σ T tg
Lr
Ei (T + tp)
r=0
satisfies
E. (0=
Proof. With constant step size Int/t. telescoping the one-step contraction gives
Clnt 4. =1 ..0 se e aa Lo + aC Lo CC+Cl =0 @(Clnt
5²0
1.
With a linearly diminishing step size a, = 4/((T + to + 1)), = O,...,t, first gives
tt the one-step contraction
1
1 L + +
to +
to + T
4C
C <
=
L
C,
2
2
ar
4
tg + F7T ++++
30

---

Published as a conference paper at ICLR 2026
Thus, the convex combination satisfies
2
10 10+T-1
l0 +
4C
4. (t 1)(t + 2tu)
TT + 1o)
Lir
L +
4
4
to +
==
(1lo + +--1+(lm) +T),, (1o + 140 11
2(6 1)(1 21o)
8C
to +
+ + (++1)(tt+ + 2te)
to 7+1
(140 - L0 - (fo + () + + 1)4+)
2(t + 1) 2to)
8C
to +
+ (1+ 1)(1 21o)
to + +1
7=0
SC1
+
24a
t²
C
=0
The convex combination removes the logarithmic dependence, and the to dependency diminishes quadratically.
L
E ANALYSIS OF CENTRAL OBJECTIVE ESTIMATION
This section directly considers central objective estimation (COE) with environment heterogeneity in Section which covers Section as a special case. We restate the learning problem in (6):
T 0 = bo.
where の0 (s)), 110
n Er= vr, and To n Ei= B. Recall that [|)()) S for all
E S. The COE algorithm is
04 0f (
(20)
where
0.b (0)) = 9t
3. (0)),
di (0)) $(s))0{ b/.
-b
The additional superscript à distinguishes the objective estimation parameters from other learning modules. We denote AA; 8 PC.
The one-step mean squared error (MSE) dynamics of (20) can be decomposed as
BI 2| = E|
2g (0)) ABT + (o 0. (0 (()))
(21)
We first analyze the cross term, then the variance term, and finally combine them to give the one-step progress. The analysis of other learning modules follows a similar pattern.
Lemma E.1 (COE descent). Ler A = Amin ( The cross term in 21 satisfies
E E(A0S aa (b(()))) > X ?
31

---

Published as a conference paper at ICLR 2026
Proof. We use the following shorthand notation: E, == B st Ej == B and Ey, B|: Ft-1), where Ft-1 is the history filtration up to time step 1. The cross term satisfies
E E ge (BS)
=EF._.
E (s) (3)),, 508 }
=1
71
"
B,o 1, (s)], A0f }
= =EFF--
Bn (
Σ
-Ep__
(T000 8 AoEi--
Note that the solution 0 satisfies 00
80 = 0. Thus,
IR/A0 u. (09))
=R5, [< - -b0) (5000- B 10;}]
=E
{ ( A0f. A
Amin (symn (b) )) A0
Lemma E.2 (COE variance). The variance term in (1 satisfies
1b19; 0.b (09)11 - 2B| 20 20
Proof. The variance term be first decomposed (11) (( where we use the fact that | The second term can be further decomposed as
El|g?
El +/4 Σ ( 81) g. (01)11 2B 21E 0. 9 (0:)
Bllg: (02)) = n 1=1
Balloi 4. (02)1 n² 2
1 E(Eiot
6.6 (02), (l (02)
H2
H1 enjoys linear variance reduction:
-112 2 E.
o²
H
6 = n² m(2G))²
The cross term H2 involves all pairs of independent local update directions. However, since each local update direction H2 is evaluated the central solution, its expectation not zero. One solution is to notice that g; 2.b Lipschitz continuous in its argument. Thus, we have | 1 (#1) ||E h.h (02) 9t affnnity-deeendet term in the variance. We adopt more "federated" approach:
> ((t) ||| = O 1]ll) = O + dag))- However, this will introduce an
H2
1
Eig. 4. (09).
Σ
((9)
n2
Eigh (8.).
E gf (00)
Eigt ((
n
Eigh
b.b (05).
Bl (89)
| (00)
n
n
00
20
| (09)1
n
=0(n 1).20
<0.
We see that by analyzing the cross terms collectively, we obtain a much tighter bound that does not depend on the affinity. Plugging the bounds of Hí and Ha back gives the desired result. I
32

---

Published as a conference paper at ICLR 2026
Combining Lemmas E.1
and E.2 with (21 gives the one-step progress of COE.
Corollary E.1 (COE one-step progress). Ler a. /4 Then, for any time step t. 20 satisfies
E A 41 2 <
(1 - A + 2(a4))
Combining Corollary
E.1 and Lemmas D.4 and D.3 gives the convergence guarantee of COE.
Corollary E.2 (COE convergence). With constant step size a lnt/(t for any time step
t> 0, 20 satisfies
E Bli0 - 8,11 = =
((x))=
0² 02 Int
With a linearly diminishing step size d = 4/T + 10 + 1)14), -= O....,4, where to > 0 ensures
that ag S /4/ (20) satisfies
o²
E 0 0-612 S 11 (1t)2nt
where Se represents the convex combination specified in Lemma
D.5
and use Jensen's inequality.
F ANALYSIS OF CENTRAL DECISION LEARNING
This section directly considers central decision learning (CDL) with environment heterogeneity and asynchronous COE (20). CDL without COE is covered as a special case with zero estimation error. We restate the learning problem (3) in Section 5
1 = Bo,
where A = variants of CDL:
1 E A(s) A and 8 Di= (s) (s). We consider two
2441 x ag
0.c (m{);
(22)
where
91
(x))
0.0
1
Σ
(((r)
(x))
Alaī bi;
22
71
0.0 or 91 (1f;
0
(1) (2)1000))) = Arri- - (())))
22
The first variant 22 corresponds to the CDL algorithm in the main text, where 9r 0.c different from the central update direction used in the personalized local learning module. the second variant 2). (s(s) = ( ( is the estimated central objective function at time step f, and we highlight this dependence by including ff in the arguments. As remarked in Appendix C. g. in the second variant (22F2) is consistent with the central update direction in the personalized local learning, and thus saves some server-side computation and communication. We will show that both variants enjoy the same convergence rate. The additional superscript distinguishes the central learning parameters from other learning modules. We denote Auf x x
The one-step MSE dynamics of can be decomposed as
E A E A 2a 0. (x)), ArT) (a (1))))
(23)
We first analyze the first variant (22-1). which is similar to the analysis of COE in Appendix E as does not involve the asynchronous COE enror.
Lemma F.I (CDL descent). Let X Amin (sym(A)) With 22 the cross in 23 satisfies
E(Ar 0. (x;) > X Ax
33

---

Published as a conference paper at ICLR 2026
Proof. Similar the proof of Lemma E.I
the cross term satisfies
Ei[A(s)of - b'(s)) - Azf )
(( 9t
0.0
(xf))
=Ba.-
i=1
= -Ey,
2A 8 (r,
2( - -b²) -AA'2. 8O) Art)
2 Amin (sym( [)|
Lemma F2 (CDL variance). With 2 the variance term in (23) satisfies
Blig 0. (1)))) <
Ar + n
Proof. Similar to the proof of Lemma H.2 the variance term can be first decomposed as
Elga (x-)) = El| Σ Aj - x + 0.0 (+ < B + 2EI., (x2)),
where the second term can be further decomposed as
Elgg 0.5 (x1))
1
E Eello + n iイj
n
AA Gb) +
Eig (x
n
<a² +0.
Combining Lemmas F.1 and F.2 with (23) gives the one-step progress of the first variant of CDL.
Corollary F.1 (CDL one-step progress). Let of S x²/(4))) Then, for any time step t, (2 satisfies
E| A 113 S(1 2/2 A 2(a
Next, analyze the second variant (22 which involves the asynchronous COE error.
Lemma F.3 (CDL +COE descent). With 2 the cross term in (23) satisfies
E(Arī, o (2f; 05))
| 2(A9)-
Proof. The cross term can be further decomposed as
Ar (xf:
Ex,
Eegh
0.0 (af; 02), Aui)
-E
E2
(x 03) - - 02))
c+i (of,
= -Ey.., (B [A( - E,, (b Azî) - Ex.., (L (S))) Aos)
=E 2A - 80, Axū) - E (Φ" A0f Art)
(24)
The first term in 24 follows a descent direction; by the definition of x
x 80, A))
=E (1A02f - 8o) (A T0), Ani)
= E AA Ar.))
Xmii (sym( (A") EE Ar
(25)
34

---

Published as a conference paper at ICLR 2026
The second term in (24 involves the estimation error from COE; by the Souchy-Schwarde inequality and Young's ss inequality,
8e
IB(*A00 Azi)l s B||A
Az
+
E
(26)
23
where Be > is constant to be determined. Plugging (25)
and (26) into (24)
gives
Pe
E Aa
0.5 (xf) 9t
x
E 5
E.
2
23
Setting 3. = X-/4 gives the desired result.
Lemma F.4 (CDL +COE variance). With (22 the variance term in (23 satisfies
El|g; 0.5 (x;; 0f) < 2 Ax 4
Ellgg
4
Proof. Similar the proof of Lemma the variance term be first decomposed
E U.c (x)² 01 El 0.0 (x 05) (2 =i) ( 05)11
where we write AN = ; 1 A and क. = E $ Therefore,
Blig (1 111 < 2G ; + 4E| A0 + 4B|lgt (xf, 03)1
Similarly, for the variance term at the stationary point, we have
1 n iT
Ellg: (x), 01)12
E 11 ci (x2, 0210
+
(22, (x
(x2, 09)
2
1=1
| x - 0000112
= 1
Plugging this back gives the desired result.
Combining Lemmas F.3 and F.4 with (23) gives the one-step progress of the second variant of CDL.
Corollary F.2 (CDL +COE one-step progress). Let af I X >²/(8)))) Then, for any time step t, 22-2) satisfies
B| - 3| -
4(ai))0''
where we use the fact that X-/G S which implies of S (X*/G>)/(8A))) & (A9)-1, and thus + (a9(X-)-1
Combining Corollaries E. and F.2 and Lemmas D.4 and D.5 gives the convergence guarantee of
CDL with asynchronous COE.
Corollary F3 (CDL convergence). With constant step size x = b6 In Int/t. ttttt,, for any time step
t t> 0, 22 satisfies
②² Int (20)2nt o² Int (AbXc)2nt
for
(2
E 1 - 12 2
for (22 with (20).
With a linearty diminishing step size a.X ensures that ai < X5 (8G2) and at < //4 (22)
a²
4/(T + to 1), t ... where to
satisfies
o² [Xo]ent
for 22
E
for (22 with (20).
(x
where 3t represents the convex combination specified in Lemma D.5
35

---

Published as a conference paper at ICLR 2026
Proof. Corollaries E.1 and F.2 fit into Lemma D.4 with 2) Pi and 2/ zz. along with at aa afxa and
C = 20
C 4 C² = 8(29)-1
Thus,
E 1 _| 166())) 1
20
640²
<(1 ar) (BiAzf 2 16(A9) + a
+
(19)2n
(1=10)21
66o²
5(1 - ar) 1 + a7
(ΛbA0)2n*
where the last inequality uses the fact that X - 101 S into Lemma D.5 gives the desired results.
Plugging the above Lyapunov function
G ANALYSIS OF PERSONALIZED COLLABORATIVE LEARNING
This section analyzes the local component of personalized collaborative learning (AffPCL), with environment heterogeneity, asynchronous COE, and asynchronous DRE that satisfies Assumption The learning problem is the most general form in (2):
A'. = B',
Vi [m].
where A' = A(8) and b' = (s). We restate the local update rule for agent
r'41 x -argi = at (o((()))) 91 #t (2t; ni (2 f()
(27)
where
(x)(()) = A(s))x -
and
nh (x) 1 n
Σ pf
Recall that &G and Ø are estimated objective and density ratio functions time step and with linear parametrization, they satisfy
(5(s) (s))0 = "s)
Thus, we highlight the dependence on estimation weights by including fj and nú in the arguments of update directions. We denote Ari x x
We first show that the local update rule follows an unbiased direction towards the local solution plus estimation emros..
Lemma G.1 (Correction). The expected local update direction satisfies
Elgl (x:) +( Aaf, 200)
where
IE(E Arf, A0;) 0 GA/AAff1 + ||001))-
36

---

Published as a conference paper at ICLR 2026
Proof. We first inspect the importance-corrected term:
Elgi
ci
(xf; 0f, ni)l
E
(()))(()) - B((62))
n Σ Eps [ [
1
-EF,-.
-EFL 1
E10 [(((s)A(s))) -
-EF._.
BRO [Io((s) (()))+(())))(AAAA
- Eny
- - + B ( An {s))0)
(28)
We notice the bias correction term exactly removes the bias in the first term above:
E,,0 ( - b((s))]]
B
(()s)
(A(s)at (((+)) n (s)ds
s (b)
1 (A(s)at ([(s))
=EF.-1
-B
E ]((+)]
ER ((s))ag B5((())
(x)))
The additional & encompasses all the estimation error:
|EEDAn Art, A0;; +]]l = An - A(s) - (z)))
<sp | + [[A(s)= - (+)
<sp((s) - - +
20 Anill + +=
where the last inequality uses Assumption | that [p[(s) - ('()) = (1). Therefore, we have
E E
Eig (=+) 15 16 (x()) (000((((())) = [=1)(()))] + IB(E(An), Ary, A0;ss));)
Corollary G.1 (AffPCL descent). Let X² Amin (sym (amm(A')) (A:)) The expected local update direction satisfies
E E Ant) 2 | - 2(2²) -BB
Proof. By Lemma G. 1 and Young's 's inequality,
E (g. Art) =Ey,_, (E (o) +
EE (E 5 (s)], Aut) + Ax))
EEF.., -B, B Ari)
=E>,,,
B¹),
Aot) + EE (E, Dai)
-E (A' Ant, Azi) (E, Az))
2A 2 - 2 7
X
B|
1
1 E1
2
Amin (A | -B1811
37

---

Published as a conference paper at ICLR 2026
For the variance, inspect the importance-corrected aggregated update direction and biased- corrected local update direction separately.
Lemma G.2 (Federated variance reduction). The variance of the importance-ccorrected aggregated update direction satisfies
Bllgg (0 0f mill < 24G +00 (n
6 262 211
Proof. Similar to (28). we decompose the importance-corrected aggregated update direction as the direction that uses the true density ratio plus estimation term:
E| El c (1f: Br n1)l &
2E
9: c
(af; 07, m/)) + 2
We can then focus on the direction with true density ratio, which again can be decomposed into local variances and covariances:
E E999 oF
(x;; B n/)11 =E
( (xf)
lp (s) (s))st a+ (=5))
1 n²
Σ jæk
Blpp (s
(s())) (s))gtt*
(xi)
H2
Different from federated variance reduction using data sampled from i.i.d. distributions, Hi also depends on how close the agents' heterogeneous environment distributions are. Suppose Fz-1 that [1 o+ (x+) & H3 for all j E [n.] Conditioned on Fi-1, we then have
H2 n
a Eys 1s))
H
j=1
2H
3
n1 +
E
pss))
n²
j=1
2H
(s(s) (2) (s)
+ E,0
n
2H
(1 (p (2))
where x2 is the chi-squared divergence. By Lemma D.3 we know that
x(p,, 1²) S max [IIp']l>, 1] || - 1 S max {Ile' lloc 13 6 We notice that the essential supremum of the density ratio has natural upper bound:
|| sup (s) sup sES ()) SS à Σ (d)
(²(s)
sup (e(s) SES 1))(())
n,
where we use the convention that 0/0- = 0. Combining the above two bounds together gives
H1 < (n + d.m))
We now bound H3. Conditioned F, have
llgi (=+) 1|/ + = / (A + 0 So GA |
Thus, we set H3 A E ||| Plugging this back gives
H1 22 E A
4E A + 40 + 5im))
(29)
<8G E |1
16 1 + 80² 80 5rm)
where the fact that -1 1 54 1.
38

---

Published as a conference paper at ICLR 2026
The covariances H2 also needs special treatment for heterogeneous environments. Unlike the homogeneous case where the local update directions are perpendicular in the sense that their covariances are zero, here the importance correction alters the geometry and requires more careful anatomy of the covariance terms. Conditioned on Fi-1, we have
H2
By (s)) 0d (ofs 8)];
Bu (s)g (2f;s))
n
j=1
Bys (p' (s)g c*
B p (5) c*k (a{;)) Bn lo (s)))))*))( (28:
n
1=1
By (p))) 0 (155 s)
R (g)) (xag
n
j=1
E (p' (s)g (xf
n
j=1
B (p' (s) ).
E 1 (s)g (x73))
2
1
Σ E (s)g (xf;s))
2
==1
Σ Basl (pell'(s00 (s)(())) -B(s))]
n
(o') ((()))||
= [m((A(s) - (E()))
= -1)))=
Note that the only inequality above omits a term of O(n-1), and thus the bound is tight when n large. Recall that of (xf) corresponds to the bias in the aggregated update direction. Thus, show that the covariance reduces nicely to the bias term, further showcasing the power of importance correction. The bias term (conditioned on Fi-1) can be further decomposed as
B = (A + x - of. + x:)) क5 (A0; + 0 - 0 + 02)11
= -I + + - xt)
(AG; + 0: 11
<4 (G Ar + (1 -! || + | + (0 01)12)
Plugging in the bounds in Items
(c) and (i) in Lemma D.I gives
| [ < 4G Ax + 4|| 1| + +20 (8 )
Removing the conditioning on F gives
H₂ < 4G F| + A0 + 3 )
(30)
Plugging 29 and (30) back gives the desired result:
B (1 < B + 40 A0 640 + 28( 2R)
where we use the fact that 8 < 8 <
We then inspect the variance of the bias-corrected local update direction.
39

---

Published as a conference paper at ICLR 2026
Lemma G.3 Affinity-based variance reduction). The variance of the bias-corrected local update direction satisfies
[x)) - ( 0E)1
A0
160 oen
where 8 min J.
Similarly, the variance term can be decomposed as the variance at the optimal solution plus the estimation error:
Proof.
Ellgi (zt) - 9i "6 (2f: 06)1
=E| =f) (b'(s)) b((s
=E | (Ar x. x --;;) $ 8
-09) |
<4G B Axt 2 =| 8E1
(estimation) (affinity)
+ 43 4 A
=I)l + 43| 0/ - 01)1
For the affinity terms, by Lemma D.2
max [E]|4] 21)1 --) < 20
Combining the above bounds gives the desired result.
We now ready prove the one-step progress of the local update in AtfPCL.
Corollary G.2 (AffPCL one-step progress). Suppose af - X/((4G)))) Then, for any time step t and
agent i, (7) satisfies
E E 2 <(1 [[0- 96(a;
+ (
5(1 )
+ + 16a ",,
+ 144(a)))0(n1 + 20 )
Proof. Combining Lemmas G.2 and G.3 gives
Elg <2( ott (xf 02 ml + E (x) g6 (x;:
E 2 64G A 96B] 1440²(n -1 + 25mm. + 4B18112
Combining the above bound with Corollary G.I gives
E Ar
1 -
2, (Øt. Az;) + o
< E | |1
28 28 5(1 - [100--003~0000000000000000144441 2
+ 9,
+ 20) 144(0))))) 25t k
Setting af < 2/ 2323322)) 2) which implies 20 5 (1) gives
E A <(1 - |
+ +=0 144( + 25
Finally, we expand 1 By Assumption
E EIE Arf A0T =E
(A MA Ao. +303)
<20 An
4 A
Plugging it back gives the desired result.
40

---

Published as a conference paper at ICLR 2026
G.1 PROOF OF THEOREM 1
Invoking Lemmas D.4 and D.5 with Corollaries E.1 F.2 and G.2 more general version of Theorem
gives us the main result. We restate
Theorem 1. We synchronize the step sizes across all learning modules by setting at of a,'² of Then, with constant step size ar, Int/(At), T O....,t, AIPPCL with various learning modules satisfies
B| - x=11== (000 (o2Int
(2))2t 8)
where
max n2 -1 3,0mm
for (27 + (22-7) + (20);
2
Been max
A
for (27 + (22
(20);
X
5 =
[Auyz
domn
(min
A n
for (27) + (22
(20) + 13
o² (Auyz
8
ocn
+ max
1o 72 for (27 + (22 (20)
+ 13
where 8cm min J.
Specifically, we highlight that AffPOL with access to the true density ratio (i.e., without 13 achieves
B - =1 Ö((rt)²t= max 2, 1),
where 6. a/\² is the agent-specific condition number, which further recovers Theorem | in the main text by noting that 8' S dem. 8' I daw + dobj. and k² S K.
On the other hand, AffPCL with DRE has worst-case complexity bounded by
Blaa 1:11 = (1n5x9)2 max max((un -,mmm)))))
where K" = o/2 and vp = max DRE.
do
1 22 x
which now depends on the conditioning of
minn{*
Proof. We only prove the first and last cases, as the other two cases follow similarly. For the last case, similar the proof of Corollary [F3 Corollaries E.1 F22 and G.2 fit into Lemma D.4 with 2 0f.
2 = oÿ, 2t = ni. and 2/ = x along with
C x C O ) C = 0(())))) C = 3 0, C+ x C C = -(2)))100)
C x C* O -1 3)
Then, the corresponding weights in Lemma D.4
w =1,
3 204.3(22 -20-00)))-06011 -000())))) 2 22 +0 (20(())))) w/ (1)) (19-1 (1-- 0 (29)
Thus, the overall MSE with constant step size In Int/tis
0² Int
E 1
2 =0
(WTY + (1:10)))
(n- + 8(m)
o2
.n -n=1) (222)) + (2ota)
+
(xxxce2 +
o²
o²
(25)2
=0
fomm (minffe, 2o, 20))
(22197) .n -n11)
(2))2t
(Ne)²
41

---

Published as a conference paper at ICLR 2026
For the first case without DRE, & O. Thus, Corollaries E.1
F.1 and G.2
fit into Lemma D.4 with
2, = fg 27
z7, and 2 = z along with
0 x C² O(o2n
C 0 kcm)
C = 0,
C x = =(a=)((a()))-
Then, the corresponding weights in Lemma a.4
w =1,
"² -20 -1(no=())1) w² =20.1 + 0= 0 (1') (1)
Thus, the overall MSE with constant step size in
Int
(min (1,, Ac)z2 n )
ago²
E||at x11 =0
(n- + 8mm)
(2#)2t
a Int
=0
mnx 8
(21)21
Similarly, by using linearly diminishing step size and convex combination of the iterates, we can remove the logarithmic factor the numerator.
42


Reviews:
[{"review_id": "RapCt3vgsP", "reviewer": "Reviewer_QeMG", "summary": "This paper introduces Personalized Collaborative Learning (PCL) for federated settings with heterogeneous agents, where each agent has unique data and goals but collaborates via a central server. The proposed algorithm, AffPCL, uses affinity-based variance reduction to balance personalization and collaboration so that agents benefit from similar peers while avoiding degradation from dissimilar ones. It achieves finite-sample convergence with a mean-squared error of order $O(t^{-1} \\\\max\\\\{n^{-1}, \\\\delta\\\\})$, adapting to agent similarity without needing prior clustering. Experiments show AffPCL outperforms independent training and matches or exceeds FedAvg across varying heterogeneity levels. Key contributions include the PCL formulation, AffPCL algorithm, convergence guarantees, and empirical validation.", "questions": "1. In the AffPCL algorithm, each agent must perform bias-corrected local updates and compute importance weights for others’ data contributions. How does the computational or communication overhead from this AffPCL algorithm affect its complexity?\n\n2. I find the authors’ argument in Lines 240–249 unclear. Could you clarify how the affinity-based variance reduction specifically benefits the proposed AffPCL algorithm?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This paper is clearly written, guiding the reader from the simplest FL algorithm to fully personalized FL with variance reduction method. The authors clearly motivate why this is needed (e.g., common global models can be suboptimal under heterogeneity).\n\n2. The proposed algorithm *AffPCL* is conceptually simple yet effective. Each agent’s update is adjusted via principled bias correction and importance weighting to account for differences in objectives and data distributions.", "weaknesses": "1. The finite-sample bounds presented in this paper appear quite loose. For instance, Prop. 1 provides an error bound for the vanilla heterogeneous FL algorithm of $\\\\tilde{O}(\\\\kappa^2 t^{-1}n^{-1})$, while Prop. 2 states the error bound for AffPCL as $\\\\tilde{O}(\\\\kappa^2 t^{-1}\\\\max\\\\{n^{-1}, \\\\delta_{obj}\\\\})$, with $\\\\delta_{obj} \\in [0,1]$ representing the effect of objective heterogeneity. Since $\\\\max\\\\{n^{-1}, \\\\delta_{obj}\\\\} \\geq n^{-1}$, the bound in Prop. 2 is no tighter than Prop. 1. The same argument applies to the bound in Thm 1.\n\n2. The experiments compare AffPCL *only* against FedAvg (which provides a single global model) and independent learning (no sharing). These are the very foundational algorithms without considering personalized collaborative learning, and the paper does not include any comparisons to existing personalized federated learning methods. There is a rich literature on approaches like per-user fine-tuning, model interpolation, or clustered federated learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "2nekdZ8ZF6", "reviewer": "Reviewer_iE7v", "summary": "This paper introduces Personalized Collaborative Learning (PCL), a unified framework for learning across heterogeneous agents that strikes a balance between personalization and inter-agent collaboration. The core algorithm, Affinity-based PCL (AffPCL), dynamically interpolates between independent and collaborative updates based on inter-agent affinity, leveraging both bias correction and importance correction mechanisms. The authors provide thorough convergence guarantees under this framework, showing that collaborative training never hurts and yields linear speedup under high affinity. Application scope spans supervised learning, reinforcement learning, and general decision-making tasks, positioning PCL as a general approach for multi-agent systems.", "questions": "How does the computational and communication complexity of AffPCL scale when deployed in realistic settings with hundreds or thousands of agents? An analysis of per-iteration runtime or messaging overhead relative to methods like FedAvg would help clarify scalability.\n\nCan the authors discuss or evaluate AffPCL on larger or more realistic datasets, beyond small-scale synthetic setups? This would help confirm whether the benefits observed in controlled environments persist under real-world data heterogeneity.\n\nThe method relies on estimating density ratios or affinity between agents. How might practitioners estimate these measures in practical federated learning scenarios where distributions are unknown or evolving?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a principled and flexible learning framework that smoothly interpolates between personalized and collaborative training, depending on the agent's similarity. By using affinity-adjusted optimization, AffPCL intelligently adapts collaboration strength instead of relying on static weighting schemes. The theoretical results are rigorous, incorporating valid assumptions and proofs that explain the model’s behavior across a range of heterogeneous settings. The clarity of exposition and strong conceptual coherence contribute to the paper’s accessibility and impact.", "weaknesses": "The limited scope of empirical evaluation undermines the practical validity of these claims. The experiments are confined to small-scale synthetic setups and do not compare AffPCL with more diverse or state-of-the-art baselines (e.g., pFedMe, Ditto, SCAFFOLD), despite some of these being mentioned in the related work. The computational and communication overhead of AffPCL, particularly in large-scale federated deployments that require density ratio estimation and per-sample corrections, remains unclear. The method’s reliance on accurately estimating density ratios or heterogeneity measures may also restrict applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "BEFvxDAP4l", "reviewer": "Reviewer_vFEm", "summary": "This paper introduces a personalized collaborative learning (PCL) framework tailored for heterogeneous agents, addressing both environmental and objective heterogeneity. The proposed method, AffPCL, is theoretically shown to reduce sample complexity by a factor of max(n^{-1},\\sigma) and analysis shows agent may obtain linear speeedup with arbitraily disimilar agents. The simulation results shows improvement over compared methods.", "questions": "Please see the weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The topic of personalized collaborative learning for heterogeneous agents is relevant and important. The formulation of PCL is a meaningful contribution, \n- The proposed framework and algorithm are clearly described and well-motivated. The formulation is broad enough to cover multiple domains such as supervised learning, reinforcement learning, and statistical decision-making.\n- Simulation results convincingly demonstrate the advantages of AffPCL, particularly under high heterogeneity settings.", "weaknesses": "- While Theorem 2 provides a lower bound, the paper would be stronger if it integrated and analyzed a practical DRE method within the AffPCL framework.\n- The quantification of objective heterogeneity is not clearly defined, which may hinder its application in real-world scenarios.\n- The current experimental evaluation is limited to synthetic linear systems. Including experiments on real-world benchmarks would significantly strengthen the empirical validity of the method.\n- Although the theoretical framework is applicable to reinforcement learning, the settings are not covered in the experiments.\n- The framework considers asynchronous importance estimation, but this aspect is not evaluated or discussed in the simulation results. \n- The proposed update rule incurs higher communication overhead compared to standard federated averaging, which may limit its scalability in certain practical settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "yAGdYcTJJS", "reviewer": "Reviewer_AoD3", "summary": "This paper proposes an affinity-based variance reduction approach that adaptively interpolates between federated learning and independent learning according to the level of objective-function heterogeneity (or data-distribution heterogeneity) among agents. Theoretically, when the data distributions across agents are homogeneous (i.e., all agents share the same objective function and the local optimal solution coincides with the global optimal solution), the proposed approach achieves the same convergence rate as federated learning. In contrast, when the data distributions are heterogeneous, the proposed approach converges to each agent's local optimal solution with a convergence rate similar to that of independent learning (where agents train in isolation). Numerical simulations validate the effectiveness of the proposed approach.\n\nMy main concerns lie in: i) overly idealized assumptions (see Weakness 1 and 2); ii) extremely heavy communication burden (see Weakness 3); iii) lack of empirical validation on benchmark datasets (see Weakness 4); iv) the reliance on a DRE oracle contradicts the \"adaptivity without prior knowledge\" claim (see Problem 1).\nWhile the theoretical analysis is sound under the given assumptions, the absence of real-world evaluations and restricted applicability of the proposed approach limit the contribution of the paper.", "questions": "See the Weaknesses section above. In addition, I have a few additional questions:\n\n1. The paper repeatedly states that the proposed approach \"does not require prior knowledge\" and \"automatically adapts to unknown heterogeneity.\" However, the algorithm relies on a density ratio term. Although the authors mention using a density-ratio estimator when this ratio is unknown, Theorem 2 shows that estimating it is nontrivial and inevitably introduces error. Therefore, the claim of \"no prior knowledge required\" is overstated unless the authors either (i) prove that the estimation is unbiased or (ii) analyze convergence in the presence of this estimated error.\n\n2. The definition of personalization in this paper differs from that in mainstream studies, as it emphasizes only local optimality while ignoring global generalization. In fact, personalized models must balance adaptation to local data with maintaining global generalization. The paper does not discuss this point, which is an omission. Please clarify this point.\n\n3. Some notation is unclear (e.g., in Equation 1, \"t\" denotes the iteration, not sample size).", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "The paper is clearly presented, and the theoretical results are well explained. The personalized collaborative learning problem studied in this work is also an active research topic.", "weaknesses": "1. **Limited applicability:** The authors assume that all agents share a common and fixed feature extractor (i.e., $\\bar{A}$; see the first formula in Section 3). This assumption is unrealistic in many real-world applications. For example, in end-to-end deep learning, both the backbone and task-specific heads are typically trained jointly. Therefore, the proposed approach is only applicable to \"frozen feature + personalized head\" settings, which limits its practical applicability.\n\n2. **Overly simplified problem setting:** The linear system in Eq. (2) essentially corresponds to the optimization problem $\\min F(x)=\\frac{1}{n}\\sum_{i=1}^{n}f_{i}(x)$, where each local objective function $f_{i}(x)=||x^{\\top}\\bar{A}x-(\\bar{b}^{i})^{\\top}x||^2$ is strongly convex and even quadratic. This formulation is considerably simpler than those used in most existing studies on federated learning or personalized collaborative learning, which typically consider convex or nonconvex objective functions. \n\n3. **Heavy communication overhead:** According to Algorithm 1 in the Appendix (which provides a summarized version of the proposed approach presented in the main text), each agent is required to send $s_{t}^{i}$, $g_{t}^{i}(x_{t}^{c})$, $g_{t}^{i,b}(\\theta_{t}^{c})$, $g_{t}^{c\\rightarrow i}(x_{t}^{c})$ to the centralized server, while the server requires to broadcast $g_{t}^{0,c}(x_{t}^{c})$, $g_{t}^{0,b}(\\theta_{t}^{c})$, $g_{t}^{c\\rightrightarrows i}(x_{t}^{c})$ to all agents. All these shared variables are high-dimensional, as they depend on the dimensionality of the underlying data in machine learning applications.\nMoreover, this bidirectional communication must occur at every iteration, which leads to an extremely heavy communication overhead. This contradicts the core motivation of FL, where communication is typically much more expensive than local computation.\n\n4. **Weakness on empirical evaluation:** The experiments are performed only on synthetic linear systems, without any evaluation on real-world datasets (e.g., FeMNIST, Shakespeare). Moreover, there lacks comparison with existing personalized FL approaches such as pFedMe, Ditto, or Clustered FL.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)