Paper deep dive

Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data

Junkang Liu, Fanhua Shang, Hongying Liu, Jin Liu, Weixin An, Yuanyuan Liu

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 127

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%

Last extracted: 7/20/2026, 8:16:57 PM

Summary

The paper identifies 'preconditioner drift' as the primary cause of instability in naive federated second-order optimizers on non-IID data, where client-side curvature geometries diverge. It proposes FedPAC, a framework that aligns local preconditioners with a global reference and corrects local updates using global curvature statistics to ensure stable convergence and improved accuracy.

Entities (9)

Federated Learning · task · 99%FedPAC · algorithm · 98%CIFAR-100 · dataset · 95%Non-IID Data · data-characteristic · 95%Preconditioner Drift · phenomenon · 95%Second-Order Optimizers · method-class · 92%Muon · optimizer · 90%SOAP · optimizer · 90%Sophia · optimizer · 90%

Relation Signals (8)

FedPAC → addresses → Preconditioner Drift

confidence 95% · To address this geometric mismatch, we propose FedPAC... explicitly decouples parameter aggregation from geometry synchronization

Muon → istypeof → Second-Order Optimizers

confidence 95% · Second-order optimizers, such as... Muon... leverage richer curvature information

Sophia → istypeof → Second-Order Optimizers

confidence 95% · Second-order optimizers, such as Sophia... leverage richer curvature information

SOAP → istypeof → Second-Order Optimizers

confidence 95% · Second-order optimizers, such as... SOAP... leverage richer curvature information

FedPAC → improves → Stability

confidence 92% · Empirically, FedPAC consistently improves stability and accuracy across vision and language tasks

Preconditioner Drift → causes → Instability

confidence 90% · naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is preconditioner drift

FedPAC → uses → Correction

confidence 90% · FedPAC explicitly decouples parameter aggregation from geometry synchronization by: (ii) Correction

FedPAC → uses → Alignment

confidence 90% · FedPAC explicitly decouples parameter aggregation from geometry synchronization by: (i) Alignment

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is \emph{preconditioner drift}: client-side second-order training induces heterogeneous \emph{curvature-defined geometries} (i.e., preconditioner coordinate systems), and server-side model averaging updates computed under incompatible metrics, corrupting the global descent direction. To address this geometric mismatch, we propose \texttt{FedPAC}, a \emph{preconditioner alignment and correction} framework for reliable federated second-order optimization. \texttt{FedPAC} explicitly decouples parameter aggregation from geometry synchronization by: (i) \textbf{Alignment} (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) \textbf{Correction} (i.e., steering local preconditioned updates using a global preconditioned direction to suppress long-term drift). We provide drift-coupled non-convex convergence guarantees with linear speedup under partial participation. Empirically, \texttt{FedPAC} consistently improves stability and accuracy across vision and language tasks, achieving up to $5.8\%$ absolute accuracy gain on CIFAR-100 with ViTs. Code is available at this https URL.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

126,128 characters extracted from source content.

Expand or collapse full text

Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data Junkang Liu Fanhua Shang∗ Hongying Liu∗ Jin Liu Weixin An Yuanyuan Liu∗ Abstract Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is preconditioner drift: client-side second-order training induces heterogeneous curvature-defined geometries (i.e., preconditioner coordinate systems), and server-side model averaging updates computed under incompatible metrics, corrupting the global descent direction. To address this geometric mismatch, we propose FedPAC, a preconditioner alignment and correction framework for reliable federated second-order optimization. FedPAC explicitly decouples parameter aggregation from geometry synchronization by: (i) Alignment (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (i) Correction (i.e., steering local preconditioned updates using a global preconditioned direction to suppress long-term drift). We provide drift-coupled non-convex convergence guarantees with linear speedup under partial participation. Empirically, FedPAC consistently improves stability and accuracy across vision and language tasks, achieving up to 5.8%5.8\% absolute accuracy gain on CIFAR-100 with ViTs. Code is available at https://anonymous.4open.science/r/FedPAC-8B24. Machine Learning, ICML shadecolorLightBlue 1 Introduction In recent years, the training paradigm of large-scale models, especially large language models (LLMs) (Devlin et al., 2019), has shifted significantly. It is increasingly recognized that the optimizer itself is part of the computational budget: under a fixed compute budget, better optimization methods can directly yield shorter training time and higher model quality. Second-order optimizers, such as Sophia (Liu et al., 2023), SOAP (Vyas et al., 2024), and Muon (Jordan et al., ), leverage richer curvature information and have demonstrated substantial acceleration in centralized large-scale training, achieving 1.5–2× faster convergence over AdamW and SGD (Abreu et al., 2025; Loshchilov et al., 2017). This evidence suggests that second-order optimization is no longer merely a theoretical luxury, but is rapidly becoming a practical backbone of large-scale deep learning. Figure 1: (a) In non-IID FL, first-order methods converge slowly, inducing little client drift. (b) Second-order methods converge faster locally and thus drift toward local optima, causing the aggregated global model to deviate from global optimum. (c) FedPAC corrects local second-order updates, yielding faster convergence and a global model closer to the global optimum. (d–f) FedPAC accelerates Sophia, Muon and SOAP to train ResNet-18 on CIFAR-100. The x-axis denotes the number of communication rounds, and the y-axis denotes test accuracy. In contrast, another highly practical and increasingly important training paradigm, Federated Learning (FL) (McMahan et al., 2017), has remained largely confined to the first-order regime in terms of optimization. Most widely adopted federated algorithms, including FedAvg (McMahan et al., 2017), and various variance-reduced variants, rely on first-order SGD updates (Bottou, 2010). Although these methods partially mitigate data heterogeneity and communication constraints, they continue to suffer from evident limitations in convergence speed, scalability to large models. This discrepancy naturally raises the following question: (a) ResNet-18, IID (b) ResNet-18, non-IID (c) ViT-Tiny, IID (d) ViT-Tiny, non-IID Figure 2: The x-axis denotes communication rounds, and the y-axis is test accuracy. (a, c): In FL on IID data, second-order optimizers converge significantly faster than SGD and AdamW for training ResNet and Transformer. (b),(d): However, on non-IID data, second-order optimizer (e.g., Local Muon) converges much more slowly and can even underperform first-order methods such as Local SGD. If second-order optimizers have already demonstrated their effectiveness in centralized large-scale training, why can we not achieve similar acceleration benefits in FL? A natural baseline is to employ Sophia/SOAP/Muon as the local optimizer on each client and aggregate parameters as in FedAvg named Local Sophia/SOAP/Muon, as shown in Figure 2. Yet, on non-IID data (Liu et al., 2025a), such a naive federated adaptation exhibits a critical failure mode: it fails to deliver the expected acceleration and often results in slower convergence or even divergence as shown in Figure 2. Our systematic experiments and mechanistic analysis indicate that this behavior is not mainly due to the computational cost of second-order information or insufficient client-side compute. Instead, it is driven by a previously underexamined phenomenon, which we refer to as preconditioner drift. In centralized training, second-order preconditioners (e.g., the matrix preconditioners implicit curvature structures exploited by SOAP and Muon) are continuously updated under a single data distribution. In federated learning, however, each client is exposed to its own non-IID data distribution. As clients perform multiple local updates, their preconditioners adapt toward the local geometry associated with their respective data distributions, which leads to the following effects as in Figure 1: (i) The preconditioners across different clients gradually drift apart in both scale and orientation; (i) The server aggregates only the model parameters via simple averaging, implicitly assuming that all clients employ similar preconditioners; (i) The resulting global update direction effectively combines inconsistent local curvature estimates to update a shared model, which severely distorts the global optimization trajectory. From an intuitive perspective, each client makes effective progress under its own curvature-adapted coordinate system; however, when these trajectories are naively aggregated, the resulting global optimization can become highly inefficient as in Figure 1. Consistent with this intuition, we observe that in non-IID data settings, directly applying second-order optimizers not only fails to replicate the acceleration achieved in centralized training, but may even perform worse than simple first-order methods as in Figure 2. This observation gives rise to an urgent and natural question: Do second-order optimizers still hold genuine potential in data heterogeneous federated learning? We answer this question affirmatively. We demonstrate that second-order optimizers can achieve substantial acceleration in FL when equipped with a dedicated framework to explicitly mitigate preconditioner drift as in Figure 1. Accordingly, we propose FedPAC (Federated Preconditioner Alignment and Correction), a principled correction and acceleration framework for second-order federated optimization. In summary, main contributions are as follows: ∙ We identify a previously underexplored failure mode of second-order federated learning: preconditioner drift—the mismatch of client-side curvature-induced geometries. We formalize it with an explicit drift metric and empirically show it strongly correlates with degraded convergence. ∙ We develop a unified framework, FedPAC, that targets this geometric mismatch by decoupling parameter updates from preconditioner synchronization: clients perform efficient local second-order updates, while preconditioners are periodically aligned and local steps are corrected using lightweight global curvature statistics. ∙ We provide drift-coupled convergence guarantees under standard smoothness and bounded-heterogeneity assumptions, where the optimization error contains an explicit drift term. We show FedPAC provably reduces this term and achieves faster convergence than first-order FL and naive second-order baselines. 2 Related Work 2.1 First-order FL under data heterogeneity In heterogeneous federated learning, most improvements of first-order methods target mitigating client drift and convergence instability caused by non-IID data. FedProx (Li et al., 2020) stabilizes local updates via a proximal regularization term, while SCAFFOLD (Karimireddy et al., 2020) theoretically analyzes client drift in FedAvg and corrects it using server–client control variates. FedCM (Xu et al., 2021) leverages client momentum to stabilize updates. Nevertheless, these methods remain fundamentally first-order, operating solely at the gradient level without explicitly leveraging second-order geometry such as the Hessian or Gauss–Newton information, leaving considerable room to explore the potential of second-order optimization in heterogeneous federated learning. Our alignment operates on the preconditioner operators that define the local metric, not on first-order states such as momentum or control variates. 2.2 Second-order optimization methods Traditional second-order methods such as Newton’s method are often prohibitively expensive. Recently, several practical second-order or quasi-second-order optimizers have been shown to scale to large models in centralized deep learning, including Shampoo (Gupta et al., 2018), SOAP (Vyas et al., 2024), Muon (Jordan et al., ), and Sophia (Liu et al., 2023). Shampoo (Gupta et al., 2018) performs structured gradient preconditioning via Kronecker-factored approximations of the Hessian. SOAP (Vyas et al., 2024) stabilizes such preconditioning by running AdamW in the induced feature basis. Muon (Jordan et al., ) orthogonalizes gradient momentum using Newton–Schulz iterations. Sophia (Liu et al., 2023) leverages lightweight diagonal Hessian estimates with element-wise clipping. 2.3 Second-order optimization in FL The utilization of second-order information in FL is still in its early stages but is developing rapidly. Fed-Sophia (Elbakary et al., 2024) adapts the Sophia optimizer to the federated setting and proposes a communication-efficient second-order federated algorithm. FedMuon (Liu et al., 2025b) introduces the matrix-orthogonalization-based Muon optimizer into FL. Our framework generalizes the alignment and correction principle in FedMuon, and thus the latter can be viewed as a special case of the former. FedPM (Ishii et al., 2025) systematically studies the empirical behavior of the second-order classical Newton method in FL, which is computational overhead and not suitable for large-scale models. 3 The Proposed Unified Algorithm 3.1 Problem Setup FL aims to optimize global model with the collaboration local clients, i.e., minimizing the following population risk: F()=1N∑i=1N(Fi():=ξi∼i[Fi(;ξi)]).F( x)= 1N _i=1^N (F_i( x):=E_ _i _i [F_i ( x; _i ) ] ). (1) The function FiF_i is the loss function on client i. ξi∼i[⋅]E_ _i _i[·] denotes conditional expectation with respect to the sample ξi _i. N is the number of clients, and x is global model. (a) Local SOAP (b) FedPAC_SOAP (c) Local SOAP (d) FedPAC_SOAP Figure 3: Performance and Preconditioner Drift of Local SOAP, FedPAC_SOAP. For SOAP, we compute the preconditioner drift in a layer-wise manner by measuring the spectral norm of the difference between each client’s left/right preconditioner and the aggregated global preconditioner. Our FedPAC_SOAP substantially reduces the preconditioner drift of Local SOAP and accelerates convergence in (a,b), reaching 40% test accuracy in fewer rounds, as shown in (c) and (d). Algorithm 1 Federated Second-Order Algorithm (FedSOA) 0: Per client, we maintain ir,k _i^r,k within a round. Hyperparameters: learning rate η, communication rounds R, the number of local updates K, the number of client N. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈ri _r in parallel do 3: ir,0← _i^r,0← 0; 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bir,k;B_i^r,k; 6: ir,k∈ℝm×n←∇Fi(ir,k;ξir,k); g_i^r,k ^m× n←∇ F_i( x_i^r,k; _i^r,k); 7: ir,k+1=UpdateState(ir,k,ir,k) _i^r,k+1=UpdateState ( _i^r,k, g_i^r,k ); 8: ~ir,k←ir,k(ir,k) g_i^r,k _ _i^r,k ( g_i^r,k ); 9: ir,k+1=ir,k−ηl~ir,k x^r,k+1_i\!=\! x^r,k_i\!\!-\! _l g_i^r,k; 10: end for 11: Δir≔ir,K−r x_i^r x_i^r,K- x^r; 12: Client i communicate (ir)( x_i^r) to Server; 13: end for 14: r+1=r+1|r|∑i=1|r|(ir,K−ir,0) x^r+1= x^r+ 1|S_r| _i=1^|S_r|( x^r,K_i- x^r,0_i); 15: end for 3.2 Unified Abstraction of Second-Order FL To investigate the mechanisms of second-order optimizers in FL, we first propose a generalized Federated Second-Order Algorithm (FedSOA), as shown in Algorithm 1. At the k-th local step of round r, client i samples a mini-batch ξir,k∼i _i^r,k _i and computes a stochastic gradient ir,k=∇Fi(ir,k;ξir,k). g_i^r,k=∇ F_i( x_i^r,k; _i^r,k). (2) Second-order optimizers maintain an internal preconditioner state ir,k _i^r,k, which parameterizes a preconditioning operator that maps the gradient to an update direction. To unify different optimizers, we define a generic operator (⋅)P_ (·) and write the local update as follows: ir,k+1=ir,k−ηlir,k(ir,k), x_i^r,k+1= x_i^r,k- _l\,P_ _i^r,k ( g_i^r,k ), (3) where ηl _l denotes the local learning rate. The specific form of P_ depends on the underlying optimizer. For example, SOAP apply structured preconditioning in a Kronecker-factored basis, Sophia rescales gradients using diagonal Hessian estimates, and Muon updates gradients in an orthogonalized rule. After each step, the preconditioner state is updated according to the optimizer rule: ir,k+1=UpdateState(ir,k,ir,k). _i^r,k+1=UpdateState ( _i^r,k, g_i^r,k ). (4) After K local steps, client i computes the model difference Δir≔ir,K−r. x_i^r x_i^r,K- x^r. A naive second-order federated baseline aggregates these updates in the same manner as FedAvg: r+1=r+∑i∈r1|r|Δir. x^r+1= x^r+ _i _r 1|S_r|\, x_i^r. (5) However, under strong statistical heterogeneity and multiple local steps, this naive aggregation leads to preconditioner drift, which significantly degrades global convergence. Three instantiations as (,)( ,P) pairs. ∙ SOAP: =L,R,(g)=A(β1m+(1−β1)g)B,A=InvRoot(L+ϵI),B=InvRoot(R+ϵI) \!=\!\L,R\,P_ (g)\!\!=\!\!A ( _1m+(1- _1)g )B,A=InvRoot(L+ε I),B=InvRoot(R+ε I). The InvRoot(⋅)InvRoot(·) denotes the inverse matrix p-th root operator. L corresponds to second-order information on the left dimension. R corresponds to second-order information on the right dimension. m denotes the momentum. ∙ Sophia: =h,(g)=clip(β1m+(1−β1)gh+ϵ,±ρ). =\h\,P_ (g)=clip\! ( _1m+(1- _1)gh+ε,\,±ρ ). h denotes the curvature (second-order information) estimate. clip(⋅)clip(·) denotes element-wise clipping. ∙ Muon: =m,(g)=Ortho(β1m+(1−β1)g). \!=\!\m\,P_ (g)=Ortho ( _1m+(1- _1)g ). Ortho(⋅)Ortho(·) denotes the orthogonalization operator. 3.3 Local Preconditioner Drift To characterize the geometric deviation of second-order optimization in federated environments, we consider an idealized centralized second-order optimizer that constructs a global preconditioner Θr ^r on the joint data distribution =1N∑i=1NiD= 1N _i=1^ND_i and performs the update r+1=r−ηΘr(∇F(r)). x^r+1= x^r-η\,P_ ^r (∇ F( x^r) ). (6) In this idealized and IID setting, the model always evolves under a unified second-order geometry characterized by r ^r. In the federated setting, each client i locally updates its preconditioner state Θir,k _i^r,k based on its own data distribution iD_i, which in turn induces a local preconditioner ir,k _i^r,k. As k and r increase, these local preconditioners progressively drift across clients. We formally define the preconditioner drift _D, as the difference gap of the preconditioner drift between the global and local preconditioners. Definition 1 (Preconditioner drift). _D of the global preconditioner r+1=1N∑i=1Nir,K ^r+1= 1N _i=1^N _i^r,K and the local preconditioners ir,Ki=1N \ _i^r,K \_i=1^N is defined as: =1N∑i=1N‖ir,K−r+1‖2. _D= 1N _i=1^NE \| _i^r,K- ^r+1 \|^2. (7) Preconditioner drift is a client-level phenomenon: different clients’ preconditioners ir,k _i^r,k gradually adapt toward their respective local geometries. As shown in Figure 3 (a,b), we observe that as data heterogeneity increases, the preconditioner drift _D becomes more severe, and the algorithm converges increasingly slowly. Algorithm 2 Federated Preconditioner Alignment and Correction Framework (FedPAC) 0: Per client, we maintain: ir,k _i^r,k within a round. Hyperparameters: learning rate η, communication rounds R, the number of local updates K, the number of client N. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈ri _r in parallel do 3: ir,0←r; _i^r,0← ^r; 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bir,kB_i^r,k; 6: ir,k∈ℝm×n←∇Fi(ir,k;ξir,k) g_i^r,k ^m× n←∇ F_i( x_i^r,k; _i^r,k); 7: ir,k+1=UpdateState(ir,k,ir,k); _i^r,k+1=UpdateState ( _i^r,k, g_i^r,k ); 8: ~ir,k←ir,k(ir,k); g_i^r,k _ _i^r,k ( g_i^r,k ); 9: ir,k+1=ir,k−ηl[(1−β)~ir,k+βGr] x^r,k+1_i\!=\! x^r,k_i\!\!-\! _l[(1\!-\!β) g_i^r,k\!+\!β g_G^r]; 10: end for 11: Δir≔ir,K−r x_i^r x_i^r,K- x^r; 12: Client i communicate (ir,ir,K)( x_i^r, _i^r,K) to Server; 13: end for 14: Gr+1=−1SKη∑i=1S(ir,K−ir,0) g_G^r+1=- 1SKη _i=1^S( x^r,K_i- x^r,0_i); 15: r+1=r+1S∑i=1S(ir,K−ir,0) x^r+1= x^r+ 1S _i=1^S( x^r,K_i- x^r,0_i); 16: r+1=1|r|∑i∈rir,K; ^r+1\;=\; 1|S_r| _i _r _i^r,K; 17: Server broadcasts (r+1,r+1,Gr+1);( x^r+1, _r+1, g_G^r+1); 18: end for 4 Our Algorithm: A Federated Preconditioner Alignment and Correction Framework To address these issues, we propose a federated preconditioner alignment and correction framework (FedPAC) for general second-order optimizers, consisting of two stages. Figure 4: An illustration of preconditioner drift, which corrects client drift through local-global alignment. (1) Aggregation of Preconditioner States (Alignment) At the end of round r, client i uploads Δir x_i^r and its local preconditioner state ir,K _i^r,K; the server then aggregates the preconditioners in Algorithm 2 (i.e., Lines 3 & 16): r+1=1|r|∑i∈rir,K. ^r+1\;=\; 1|S_r| _i _r _i^r,K. (8) At the beginning of round r+1r\!+\!1, the server broadcasts (r,r)( x^r, ^r) to the participating clients. Upon receiving them, client i aligns its local preconditioner state with the global reference r ^r to obtain the aligned preconditioner used for the new round: ir,0←r _i^r,0← ^r. The client then uses ir,0 _i^r,0 as the initial preconditioner state for round r+1r\!+\!1 and performs the local second-order updates according to (3). This mechanism preserves local second-order adaptivity while periodically pulling the optimization geometry back toward a shared reference. (2) Local Preconditioner Correction (Correction) On each client, we correct the preconditioner drift, as follows: ir,k+1=ir,k−ηl[(1−β)~ir,k+βGr], x^r,k+1_i\!=\! x^r,k_i\!\!-\! _l[(1\!-\!β) g_i^r,k\!+\!β g_G^r], (9) where Gr=−1SKη∑i=1S(ir−1,K−ir−1,0) g_G^r\!=\!- 1SKη _i=1^S ( x_i^r-1,K\!-\! x_i^r-1,0 ) is the estimated global update, obtained by averaging the preconditioned-gradient global change from the previous round. β is the trade-off coefficient between local and global updates in Figure 4 and Algorithm 2 (Lines 9 & 14). Overall, the above steps constitute a unified correction and acceleration framework for second-order federated optimization. The underlying local optimizer can be any structured second-order method, such as Sophia, Muon, SOAP. In this work, we instantiate our framework with three representative second-order optimizers, yielding FedPAC_Sophia, FedPAC_Muon, and FedPAC_SOAP, respectively. (a) CIFAR-100,Dir-0.1 (b) CIFAR-100,Dir-0.05 (c) Tiny-ImageNet,Dir-0.1 (d) Tiny-ImageNet,Dir-0.05 (e) CIFAR-100,Dir-0.1 (f) CIFAR-100,Dir-0.05 (g) Tiny-ImageNet,Dir-0.1 (h) Tiny-ImageNet,Dir-0.05 Figure 5: Test accuracy versus communication rounds. The x-axis denotes communication rounds and the y-axis denotes the test accuracy. Panels (a–d) correspond to training with ResNet-18, while panels (e–h) correspond to training with ViT-Tiny. Table 1: Test accuracy of each method on CIFAR-100 and Tiny-Imagenet using ResNet-18 and ViT-Tiny over 300 communication rounds under Dir-0.1 and Dir-0.05 (100 clients, 10% participation, batch size 50, K=50K=50). Method ResNet-18 ViT-Tiny CIFAR-100 Tiny-Imagenet CIFAR-100 Tiny-Imagenet Dir-0.1 Dir-0.05 Dir-0.1 Dir-0.05 Dir-0.1 Dir-0.05 Dir-0.1 Dir-0.05 FedAvg 60.17 56.75 47.48 43.80 27.24 23.42 15.68 14.05 SCAFFOLD 60.69 56.43 47.76 43.92 26.86 23.23 15.70 14.21 FedCM 66.61 62.65 41.16 36.00 16.95 14.74 8.88 8.15 Local AdamW 59.23 55.24 44.01 40.00 37.57 36.06 24.31 21.35 Local Sophia 56.65 50.89 41.23 36.15 34.05 32.25 22.49 21.14 FedPAC_Sophia 59.96 (↑3.31 3.31) 53.66 (↑2.77 2.77) 43.81 (↑2.58 2.58) 36.37 (↑0.22 0.22) 39.79 (↑5.74 5.74) 32.71 (↑0.46 0.46) 23.01 (↑0.52 0.52) 22.37 (↑1.23 1.23) Local Muon 67.26 49.86 52.83 34.76 44.00 39.68 30.51 28.25 FedPAC_Muon 71.85 (↑4.59 4.59) 65.56 (↑15.70 15.70) 57.95 (↑5.12 5.12) 54.00 (↑19.24 19.24) 47.81 (↑3.81 3.81) 41.76 (↑2.08 2.08) 31.45 (↑0.94 0.94) 30.25 (↑2.00 2.00) Local SOAP 68.44 58.16 54.42 50.02 49.41 41.68 33.30 30.36 FedPAC_SOAP 69.25 (↑0.81 0.81) 64.16 (↑6.00 6.00) 55.62 (↑1.20 1.20) 51.81 (↑1.79 1.79) 51.16 (↑1.75 1.75) 47.55 (↑5.87 5.87) 34.32 (↑1.02 1.02) 31.33 (↑0.97 0.97) 5 Theoretical Analysis In this section, we give the convergence theoretical analysis of our proposed FedPAC algorithm. Firstly, we state some standard assumptions for the non-convex function f. Assumption 5.1 (Smoothness). The non-convex fif_i is a L-smooth function for all i∈[N]i∈[N], i.e., ‖∇fi()−∇fi()‖≤L‖−‖\|∇ f_i( x)-∇ f_i( y)\|≤ L\| x- y\|, for all ,∈ℝd x, y ^d. Assumption 5.2 (Bounded Stochastic Gradient). ir=∇fi(ir,ξir) g_i^r=∇ f_i( x_i^r, _i^r) computed by using a sampled mini-batch data ξir _i^r in the local client i is an unbiased estimator of ∇fi∇ f_i with bounded variance, i.e., ξir[ir]=∇fi(ir)E_ _i^r[ g_i^r]=∇ f_i( x_i^r) and ξir‖ir−∇fi(ir)‖2≤σl2E_ _i^r\| g_i^r-∇ f_i( x_i^r)\|^2≤ _l^2, for all ir∈ℝd x_i^r ^d. Assumption 5.3 (Bounded Heterogeneity). The dissimilarity between local clients is bounded on the gradients, i.e., ‖∇fi()−∇f()‖2≤σg2\|∇ f_i( x)-∇ f( x)\|^2≤ _g^2, for all ∈ℝd x ^d. Assumption 5.4 (Preconditioner Coercivity and Boundedness). For any preconditioner state Θ and any vector ∈ℝd v ^d, the preconditioned mapping PΘ(⋅)P_ (·) satisfies the following two properties: (i) (coercivity) ⟨,PΘ()⟩≥μ‖2 v,P_ ( v) ≥μ\| v\|^2 for some constant μ>0μ>0; (i) (boundedness) ‖PΘ()‖≤M‖\|P_ ( v)\|≤ M\| v\| for some constant M>0M>0. Assumption 5.5 (Lipschitz Continuity in Preconditioner State). There exists a constant L>0L_ >0 such that for any two preconditioner states ,′ , and any vector ∈ℝd v ^d, ‖P()−P′()‖≤L‖−′‖⋅‖.\|P_ ( v)-P_ ( v)\|≤ L_ \| - \|·\| v\|. Our convergence theory holds for the class of state-dependent preconditioned mappings characterized by Assumptions 5.4 and 5.5 (spectrally bounded and Lipschitz in the preconditioner state), which is reasonably satisfied by SOAP/Muon/Sophia. Theorem 5.6 (Convergence of FedSOA for non-convex functions ). Under Assumptions 5.1, 5.2, 5.3, 5.4, 5.5, ∃G2>0s.t.supr,i,k‖gir,k‖2≤G2,∃\,G^2>0\ s.t.\ _r,i,kE\|g_i^r,k\|^2≤ G^2, if we take g0=0g^0=0 and choose the stepsize η such that η≤min⁡μ4LM2K,12LMK,η\;≤\; \ μ4LM^2K,\ 12LMK \, then FedSOA satisfies 1R∑r=0R−1‖∇f(r)‖2≲LΔR+LΔR⋅σl2+σg2+SκΘΔ¯DSK, 1R _r=0^R-1E \|∇ f( x^r) \|^2\! \! L R+ L R\!·\! _l^2+ _g^2+S _ \, _DSK, where Δ=f(0)−f⋆ =f( x^0)-f , and Δ¯D=1R∑r=0R−1ΔDr,ΔDr=1S∑i∈Sr‖Θir,K−Θ¯r,K‖2. _D= 1R _r=0^R-1 _D^r, _D^r= 1S _i∈ S_rE \| _i^r,K- ^r,K \|^2. and the drift-to-noise coupling constant is κΘ:=LΘ2G2. _ \;:=\;L_ ^2\,G^2. Here S is the number of participating clients per round, K is the number of local iterations, and R is the total number of communication rounds. The notation ≲ hides absolute constants depending only on (L,μ,M)(L,μ,M). Theorem 5.7 (Convergence of FedPAC for non-convex functions). Under Assumptions 5.1, 5.2, 5.4, and 5.5, and the following second-moment bound ∃G2>0s.t.supr,i,k‖gir,k‖2≤G2,∃\,G^2>0\ s.t.\ _r,i,kE\|g_i^r,k\|^2≤ G^2, if we take g0=0g^0=0 and choose parameters such that η≤min⁡μ4LM2K,12LMK,η\;≤\; \ μ4LM^2K,\ 12LMK \, then FedPAC satisfies 1R∑r=0R−1‖∇f(r)‖2≲LΔR+LΔσl2+κΘ⋅Δ¯DSKR. 1R _r=0^R-1E \|∇ f( x^r) \|^2\; \; L R+ L \, _l^2+ _ · _DSKR. (1) What limits second-order FL? Theorem 5.6 establishes the non-convex convergence of FedSOA and, more importantly, makes the preconditioner drift explicit: beyond the standard optimization term LΔR L R and the stochastic/heterogeneity noise σl2+σg2SK _l^2+ _g^2SK, an additional penalty SκΘΔ¯DS _ _D enters the effective noise. This shows that inconsistent client-side preconditioner states can dominate convergence even when gradients are well-estimated. (2) How does FedPAC close the gap? Theorem 5.7 shows that FedPAC eliminates the explicit heterogeneity term σg2 _g^2 by aligning/correcting preconditioner dynamics, leaving only the drift-dependent remainder Δ¯D _D. When Δ¯D _D is small, FedPAC achieves near-ideal scaling dominated by σl2/(SK) _l^2/(SK), making drift measurable and controllable in practice. 6 Experiments Datasets. We evaluate FedPAC on both vision and language tasks. (i) For image classification, we use CIFAR-100 (Krizhevsky et al., 2009), and Tiny-ImageNet (Le and Yang, 2015). (i) For NLP tasks, we adopt C4 (Raffel et al., 2020) dataset. To simulate data heterogeneity across clients, we follow the Dirichlet partitioning scheme (Hsu et al., 2019). For Dir-α partitioning, smaller α indicates more severe data heterogeneity. Model Architectures. We explore a variety of model types: (i) ResNet-18 (He et al., 2016) as a representative convolutional neural network (CNN), (i) Vision Transformer (ViT-Base) and ViT-Tiny (Dosovitskiy et al., 2020) for Vision Transformers, and (i) LLaMA (Touvron et al., 2023) for large-scale language model. Baselines. We compare our method against state-of-the-art FL algorithms: FedAvg (Local SGD) (McMahan et al., 2017), SCAFFOLD (Karimireddy et al., 2020), FedCM (Xu et al., 2021), Local AdamW, Local Sophia, Local Muon and Local SOAP. Under our FedPAC framework, we instantiate three variants: FedPAC_Sophia, FedPAC_Muon, and FedPAC_SOAP. In the Appendix (Table), we compare additional FL algorithms designed to address data heterogeneity. Hyperparameter Settings. For FedAvg, SCAFFOLD, FedCM, the lrlr is selected from 10−2, 3×10−2, 5×10−2, 10−1, 3×10−1\10^-2,\ 3× 10^-2,\ 5× 10^-2,\ 10^-1,\ 3× 10^-1\, and choose the best value lr=10−1lr=10^-1. with a weight decay of 0.0010.001. For Local AdamW, the lrlr is selected from 10−4, 3×10−4, 5×10−4, 8×10−4, 10−3\10^-4,\ 3× 10^-4,\ 5× 10^-4,\ 8× 10^-4,\ 10^-3\, and choose the best value lr=3×10−4lr=3\!×\!10^-4. We set the learning rates for Local Sophia, Local Muon, and Local SOAP to 3×10−43\!×\!10^-4, 3×10−23\!×\!10^-2, and 3×10−33\!×\!10^-3, respectively. Their weight decay is 0.010.01. We apply cosine learning rate decay, and set FedPAC to =0.5 β\!=\!0.5, weight decay 0.010.01. We set the learning rate of FedPAC variants to be same with Local Sophia, Local Muon and Local SOAP. Additional hyperparameter configurations are detailed in the Appendix. We release all code, configuration files to ensure full reproducibility. All results are averaged over 5 runs with std reported with seeds 42, 43, 44, 45, 46. All experiments were performed on a single NVIDIA RTX 4090 GPU. Table 2: Comparison of test accuracy and training loss for Vision Transformer (ViT-Base) under Dir-0.1 with 100 rounds (50 clients, 10% participation, batch size 16, K=50K=50). Method CIFAR-100 Tiny-ImageNet Test Acc Loss Test Acc Loss FedAvg 89.2889.28 0.332 86.9786.97 0.421 SCAFFOLD 89.3589.35 0.328 86.8586.85 0.432 FedCM 87.4487.44 0.592 84.2484.24 0.708 Local AdamW 90.1090.10 0.281 86.6586.65 0.252 Local Sophia 90.2190.21 0.275 86.1686.16 0.188 FedPAC_Sophia 90.3890.38 0.262 86.5586.55 0.175 Local Muon 90.3290.32 0.189 87.2287.22 0.184 FedPAC_Muon 90.4690.46 0.165 87.6387.63 0.171 Local SOAP 90.4490.44 0.151 87.7287.72 0.212 FedPAC_SOAP 90.65 0.129 87.92 0.156 Table 3: The train loss of each method on C4 data using LLaMA 60M, LLaMA 130M, LLaMA 350M over 100 communication rounds (20 clients, 20% participation, batch size 16, K=50K=50). Method C4 (Train Loss) LLaMA 60M LLaMA 130M LLaMA 350M FedAvg 4.156 4.258 4.354 SCAFFOLD 4.028 4.238 4.365 FedCM 4.231 4.356 4.426 Local AdamW 3.213 3.418 3.519 Local Sophia 3.201 3.389 3.521 FedPAC_Sophia 3.188 3.367 3.485 Local Muon 3.156 3.256 3.365 FedPAC_Muon 3.145 3.229 3.352 Local SOAP 3.148 3.245 3.312 FedPAC_SOAP 3.121 3.215 3.341 Table 4: Impact of β on FedPAC_SOAP with ViT-Tiny and ResNet-18 on CIFAR-100 under Dir-0.05. Model, Dataset, β 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 ResNet-18, CIFAR-100 62.05 62.35 63.10 63.57 63.75 64.16 63.75 61.75 61.70 58.86 ViT-Tiny, CIFAR-100 44.56 45.28 45.63 46.22 46.84 47.55 46.86 46.59 46.05 45.66 6.1 Results on Convolutional Neural Networks Training on CIFAR-100 with ResNet-18. Table 1 and Figure 1 present the test accuracy on CIFAR-100 and Tiny-ImageNet using ResNet-18. Accuracy under heterogeneity. Across both datasets, FedPAC consistently improves the corresponding local second-order optimizer, with the largest gains in highly non-IID regimes where preconditioner drift is severe. For example on CIFAR-100, Local Muon drops from 67.26%67.26\% (Dir-0.1) to 49.86%49.86\% (Dir-0.05), whereas FedPAC_Muon achieves 71.85%71.85\% and 65.56%65.56\%, respectively. FedPAC_SOAP also strengthens SOAP on CIFAR-100, improving from 68.44%68.44\%/58.16%58.16\% (Local SOAP) to 69.25%69.25\%/64.16%64.16\% under Dir-0.1/Dir-0.05. Similar trends hold on Tiny-ImageNet: Local Muon attains 52.83%52.83\%/34.76%34.76\% (Dir-0.1/Dir-0.05), while FedPAC_Muon reaches 57.95%57.95\%/54.00%54.00\%. Improvements for Sophia are also consistent (e.g., CIFAR-100: 56.65%→59.96%56.65\%→ 59.96\% at Dir-0.1 and 50.89%→53.66%50.89\%→ 53.66\% at Dir-0.05). Stability and drift mitigation. As illustrated in Figure 5, FedPAC reduces preconditioner drift and yields faster, more stable convergence under strong heterogeneity, explaining why aggressive second-order methods (e.g., Muon/SOAP) remain effective in federated CNN training. 6.2 Results on Vision Transformer Training ViT-Tiny from scratch. Results are reported in Table 1. FedPAC yields the strongest performance among second-order variants, especially under non-IID data. On CIFAR-100, FedPAC_SOAP improves over Local_SOAP from 49.41%/41.68%49.41\%\!/41.68\% to 51.16%/47.55%51.16\%\!/47.55\% under Dir-0.1/0.050.1/0.05, respectively. On Tiny-ImageNet, FedPAC_SOAP also achieves the best overall accuracy (e.g. 31.33%31.33\% at Dir-0.050.05), while FedPAC_Muon shows clear robustness gains under heterogeneity (Dir-0.1/0.050.1/0.05: 31.45%/30.25%31.45\%/30.25\% vs. 30.51%/28.25%30.51\%/28.25\% for Local_Muon). Overall, these results indicate that FedPAC effectively stabilizes federated ViT optimization and preserves accuracy as heterogeneity increases. Pretrained ViT-Base. Since small-scale datasets can limit ViT performance when trained from scratch, we additionally evaluate a pretrained ViT-Base under Dir-0.10.1 (Table 2). FedPAC continues to provide consistent improvements: FedPAC_SOAP reaches 90.65%90.65\% on CIFAR-100 and 87.92%87.92\% on Tiny-ImageNet, surpassing the corresponding local second-order baselines (e.g., Local_SOAP: 90.44%/87.72%90.44\%/87.72\%). These results confirm the effectiveness of FedPAC on federated vision Transformers across both from-scratch and pretrained settings. 6.3 C4 federated pre-training with LLaMA. Table 3 summarizes C4 federated pre-training results after 100 rounds (20 clients, 20% participation) for LLaMA-60M/130M/350M. Standard FL optimizers struggle in this heterogeneous, partial-participation regime (e.g., FedAvg: 4.156/4.258/4.354), whereas locally using modern adaptive/second-order optimizers already yields large gains (e.g., Local SOAP: 3.148/3.245/3.312). FedPAC consistently matches or improves the strongest local baselines across scales, with the best results achieved by FedPAC_SOAP on 60M/130M (3.121/3.215) and competitive performance on 350M. Compared to FedAvg, the best configuration reduces training loss by ∼ 1.0 absolute (e.g., 4.156→ 3.121 on 60M), highlighting the practical benefit of stabilizing second-order methods in FL. 6.4 Ablation Study A1: Sensitivity to β (correction strength). Table 4 evaluates FedPAC_SOAP under Dir-0.05 on CIFAR-100. Performance is robust for moderate β and peaks around β=0.5β=0.5 for both backbones (ResNet-18: 64.1664.16, ViT-Tiny: 47.5547.55), indicating that a balanced global-direction correction is most effective. Too small β under-utilizes the correction signal (e.g., β=0β=0), while overly large β over-regularizes toward the global direction and degrades accuracy (notably on ResNet-18 for β≥0.7β≥ 0.7). We thus use β=0.5β=0.5 as the default in all experiments. A2: Component-wise ablation (Alignment vs. Correction). Table 5 shows that both modules are necessary. On Dir-0.05 CIFAR-100, Alignment-only and Correction-only each improves over Local SOAP, but Full FedPAC_SOAP achieves the best accuracy, indicating that alignment and correction are complementary. Table 5: Ablation on FedPAC_SOAP components under Dir-0.05 on CIFAR-100. We report test accuracy after 300 rounds. Variant ResNet-18 ViT-Tiny Local SOAP 58.1658.16 41.6841.68 w/o preconditioner alignment (Θ¯ ) 61.1261.12 43.6743.67 w/o preconditioner correction (ΔG _G) 62.0562.05 44.5644.56 FedPAC_SOAP (full) 64.1664.16 47.5547.55 A3: Communication-efficient preconditioner aggregation. Table 6 studies a compressed variant, FedPAC_SOAP_light, which uploads the SOAP preconditioner using SVD compression. While full FedPAC_SOAP improves accuracy over Local SOAP (47.55 vs. 41.68) at the cost of 3×3× communication (68.4 MB/round), FedPAC_SOAP_light largely preserves the gain (46.75) with near-local bandwidth (25.1 MB/round, 1.1×1.1×). Notably, the computation overhead remains small across variants (5.56–5.91 s/round), indicating that SVD-based compression offers a favorable accuracy–communication trade-off. Table 6: Communication-efficient preconditioner aggregation for SOAP on CIFAR-100, ViT-Tiny under Dir-0.05. CommCost denotes communication cost per round (MB), and CompCost denotes computation time per round (s). Aggregation Acc CommCost CompCost Local SOAP 41.68 22.8 MB (1×) 5.56 s FedPAC_SOAP 47.55 68.4 MB (3×) 5.85 s FedPAC_SOAP_light 46.75 25.1 MB (1.1×) 5.91s 7 Conclusion We study how modern second-order optimizers behave in federated learning under non-IID data. We identify preconditioner drift as a key source of instability: after multiple local steps, clients adapt preconditioners to different local geometries, and naive aggregation across mismatched geometries distorts global updates. To mitigate this issue, we propose FedPAC, a preconditioner alignment and correction framework that aggregates and broadcasts a global reference preconditioner to align local geometry and suppress drift accumulation. We provide convergence guarantees for non-convex objectives and demonstrate consistent improvements in stability and accuracy on heterogeneous vision benchmarks across SOAP, Sophia, and Muon. 8 Limitations and Future Work Communication. Synchronizing preconditioner states increases communication, especially for large layers. Future work will explore efficient transmission/aggregation via low-rank compression. IID regimes. When data are close to IID, the alignment/correction may bring limited benefit and can underperform the original optimizer. Computation. Second-order local updates add computation overhead (about 1.2×1.2×–1.5×1.5× vs. first-order), which may hinder deployment on constrained clients. More efficient approximations and system optimizations are needed. Impact Statement This work proposes FedPAC, a federated optimization framework that improves the stability and efficiency of curvature-aware (preconditioned/second-order) training under client heterogeneity and partial participation. More reliable and communication-efficient federated training can lower resource and energy costs and broaden access to collaborative learning in domains where data cannot be centralized (e.g., on-device personalization or cross-silo settings). However, stronger federated training may also amplify risks from inappropriate deployment, including unfair decisions, privacy leakage, or misuse. Federated learning does not by itself guarantee privacy or security, and our method provides no formal privacy guarantees; practitioners should combine it with established protections (e.g., secure aggregation, differential privacy, and auditing) and evaluate robustness and fairness prior to deployment. References N. Abreu, N. Vyas, S. Kakade, and D. Morwani (2025) The potential of second-order optimization for llms: a study with full gauss-newton. arXiv preprint arXiv:2510.09378. Cited by: §1. L. Bottou (2010) Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT’2010: 19th International Conference on Computational StatisticsParis France, August 22-27, 2010 Keynote, Invited and Contributed Papers, p. 177–186. Cited by: §1. J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019) BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL, Cited by: §1. A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. (2020) An image is worth 16x16 words: transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. Cited by: §D.4, §D.5, §6. A. Elbakary, C. B. Issaid, M. Shehab, K. Seddik, T. ElBatt, and M. Bennis (2024) Fed-sophia: a communication-efficient second-order federated learning algorithm. In ICC 2024-IEEE International Conference on Communications, p. 950–955. Cited by: §2.3. V. Gupta, T. Koren, and Y. Singer (2018) Shampoo: preconditioned stochastic tensor optimization. In International Conference on Machine Learning, p. 1842–1850. Cited by: §2.2. K. He, X. Zhang, S. Ren, and J. Sun (2016) Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, p. 770–778. Cited by: §D.3, §6. T. H. Hsu, H. Qi, and M. Brown (2019) Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335. Cited by: §6. H. Ishii, K. Niwa, H. Sawada, A. Fujino, N. Harada, and R. Yokota (2025) FedPM: federated learning using second-order optimization with preconditioned mixing of local parameters. arXiv preprint arXiv:2511.09100. Cited by: §2.3. [10] K. Jordan, Y. Jin, V. Boza, Y. Jiacheng, F. Cecista, L. Newhouse, and J. Bernstein Muon: an optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan. github. io/posts/muon 6. Cited by: §1, §2.2. S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh (2020) Scaffold: stochastic controlled averaging for federated learning. In International Conference on Machine Learning, p. 5132–5143. Cited by: §2.1, §6. A. Krizhevsky, G. Hinton, et al. (2009) Learning multiple layers of features from tiny images. Cited by: 1st item, §6. Y. Le and X. Yang (2015) Tiny imagenet visual recognition challenge. CS 231N 7 (7), p. 3. Cited by: §6. T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith (2020) Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, Cited by: §2.1. H. Liu, Z. Li, D. Hall, P. Liang, and T. Ma (2023) Sophia: a scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342. Cited by: §1, §2.2. J. Liu, Y. Liu, F. Shang, H. Liu, J. Liu, and W. Feng (2025a) Improving generalization in federated learning with highly heterogeneous data via momentum-based stochastic controlled weight averaging. In Forty-second International Conference on Machine Learning, Cited by: §1. J. Liu, F. Shang, J. Zhou, H. Liu, Y. Liu, and J. Liu (2025b) FedMuon: accelerating federated learning with matrix orthogonalization. External Links: 2510.27403, Link Cited by: §2.3. Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo (2021) Swin transformer: hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, p. 10012–10022. Cited by: §D.5. I. Loshchilov, F. Hutter, et al. (2017) Fixing weight decay regularization in adam. arXiv preprint arXiv:1711.05101 5 (5), p. 5. Cited by: §1. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas (2017) Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, p. 1273–1282. Cited by: §1, §6. C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu (2020) Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research 21 (140), p. 1–67. Cited by: §6. H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. (2023) Llama: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: §6. N. Vyas, D. Morwani, R. Zhao, M. Kwun, I. Shapira, D. Brandfonbrener, L. Janson, and S. Kakade (2024) Soap: improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321. Cited by: §1, §2.2. J. Xu, S. Wang, L. Wang, and A. C. Yao (2021) Fedcm: federated learning with client-level momentum. arXiv preprint arXiv:2106.10874. Cited by: §2.1, §6. Appendix A Appendix A: Proof of Theorem 1 and Convergence Analysis A.1 FedPAC Algorithm A.2 Assumption We analyze generalization based on following assumptions: Assumption A.1. (Smoothness). FiF_i is L-smooth for all i∈i∈ [N][N], ‖∇Fi(1)−∇Fi(2)‖≤L‖1−2‖ \|∇ F_i( x_1)-∇ F_i( x_2) \|≤ L\| x_1- x_2\| (10) for all 1,2 x_1, x_2 in its domain and i∈[N]i∈[N]. Assumption A.2. (Bounded variance of data heterogeneity). The global variability of the local gradient of the loss function is bounded by σg2 _g^2 for all i∈[N]i∈[N], ‖∇Fi()−∇F()‖2≤σg2 \|∇ F_i ( x )-∇ F ( x ) \|^2≤ _g^2 (11) Assumption A.3. (Bounded variance of stochastic gradient). The stochastic gradient ∇Fi(,ξi)∇ F_i ( x, _i ), computed by the i-th client of model parameter x using mini-batch ξi _i, is an unbiased estimator of ∇Fi()∇ F_i( x) with variance bounded by σl2 _l^2, i.e., ξi‖∇Fi(,ξi)−∇Fi()‖2≤σl2E_ _i \|∇ F_i ( x, _i )-∇ F_i( x) \|^2≤ _l^2 (12) for all i∈[N]i∈[N], where the expectation is over all local datasets. Assumption A.4 (Preconditioner Coercivity and Boundedness). For any preconditioner state Θ and any vector ∈ℝd v ^d, the preconditioned mapping PΘ(⋅)P_ (·) satisfies the following two properties: (i) (coercivity) ⟨,PΘ()⟩≥μ‖2 v,P_ ( v) ≥μ\| v\|^2 for some constant μ>0μ>0; (i) (boundedness) ‖PΘ()‖≤M‖\|P_ ( v)\|≤ M\| v\| for some constant M>0M>0. Assumption A.5 (Lipschitz Continuity in Preconditioner State). There exists a constant LΘ>0L_ >0 such that for any two preconditioner states Θ,Θ′ , and any vector ∈ℝd v ^d, ‖PΘ()−PΘ′()‖≤LΘ‖Θ−Θ′‖⋅‖.\|P_ ( v)-P_ ( v)\|≤ L_ \| - \|·\| v\|. Appendix B Proof of Drift-Coupled Convergence for FedSOA and Alignment to Theorem 5.6 B.1 Setup and Notation Let F(x):=1N∑i=1NFi(x)F(x):= 1N _i=1^NF_i(x) be a (possibly non-convex) objective. In communication round r, a set SrS_r of S clients participates. Each client performs K local updates: xir,k+1=xir,k−ηg~ir,k,g~ir,k:=PΘir,k(gir,k),gir,k:=∇Fi(xir,k;ξir,k).x_i^r,k+1=x_i^r,k-η g_i^r,k, g_i^r,k:=P_ _i^r,k\! (g_i^r,k ), g_i^r,k:=∇ F_i(x_i^r,k; _i^r,k). (13) The server aggregates: xr+1=xr+1S∑i∈Sr(xir,K−xir,0)=xr−ηUr,Ur:=1S∑i∈Sr∑k=0K−1g~ir,k.x^r+1=x^r+ 1S _i∈ S_r (x_i^r,K-x_i^r,0 )=x^r-η\,U^r, U^r:= 1S _i∈ S_r _k=0^K-1 g_i^r,k. (14) Define the (per-round) averaged preconditioner state Θ¯r,K:=1S∑i∈SrΘir,K ^r,K:= 1S _i∈ S_r _i^r,K and drift metric ΔDr:=1S∑i∈Sr‖Θir,K−Θ¯r,K‖2. _D^r:= 1S _i∈ S_rE \| _i^r,K- ^r,K \|^2. (15) B.2 Assumptions (A1) Smoothness. Each FiF_i is L-smooth: ‖∇Fi(x)−∇Fi(y)‖≤L‖x−y‖\|∇ F_i(x)-∇ F_i(y)\|≤ L\|x-y\|. (A2) Unbiased stochastic gradients with bounded variance. [gir,k∣xir,k]=∇Fi(xir,k)E[g_i^r,k x_i^r,k]=∇ F_i(x_i^r,k) and ‖gir,k−∇Fi(xir,k)‖2≤σl2E\|g_i^r,k-∇ F_i(x_i^r,k)\|^2≤ _l^2. (A3) Bounded heterogeneity (FedSOA only). ‖∇Fi(x)−∇F(x)‖2≤σg2\|∇ F_i(x)-∇ F(x)\|^2≤ _g^2 for all i,xi,x. (P1) Preconditioner coercivity and boundedness. There exist constants 0<μ≤M0<μ≤ M such that for all Θ and all v, ⟨v,PΘ(v)⟩≥μ‖v‖2,‖PΘ(v)‖≤M‖v‖. v,P_ (v) ≥μ\|v\|^2, \|P_ (v)\|≤ M\|v\|. (16) (P2) Lipschitz continuity in the preconditioner state. There exists LΘ>0L_ >0 such that for all Θ,Θ′ , and all v, ‖PΘ(v)−PΘ′(v)‖≤LΘ‖Θ−Θ′‖⋅‖v‖.\|P_ (v)-P_ (v)\|≤ L_ \| - \|·\|v\|. (17) (A4) Second-moment bound. There exists G2G^2 such that supr,i,k‖gir,k‖2≤G2 _r,i,kE\|g_i^r,k\|^2≤ G^2. B.3 A Key Lemma: Injecting Drift into the Upper Bound Lemma B.1 (Drift-induced preconditioned disagreement). Let Θ¯ be any reference state and Θii∈Sr\ _i\_i∈ S_r be local states. Then for any vectors vii∈Sr\v_i\_i∈ S_r, ‖1S∑i∈Sr(PΘi(vi)−PΘ¯(vi))‖2≤LΘ2(1S∑i∈Sr‖vi‖2)(1S∑i∈Sr‖Θi−Θ¯‖2). \| 1S _i∈ S_r (P_ _i(v_i)-P_ (v_i) ) \|^2≤ L_ ^2 ( 1S _i∈ S_r\|v_i\|^2 ) ( 1S _i∈ S_r\| _i- \|^2 ). (18) Proof. By (17), ‖1S∑i(PΘi(vi)−PΘ¯(vi))‖≤1S∑iLΘ‖Θi−Θ¯‖⋅‖vi‖. \| 1S _i(P_ _i(v_i)-P_ (v_i)) \|≤ 1S _iL_ \| _i- \|·\|v_i\|. Apply Cauchy–Schwarz: (1S∑iaibi)2≤(1S∑iai2)(1S∑ibi2), ( 1S _ia_ib_i )^2≤ ( 1S _ia_i^2 ) ( 1S _ib_i^2 ), with ai=‖Θi−Θ¯‖a_i=\| _i- \| and bi=‖vi‖b_i=\|v_i\|. ∎ Instantiation with ΔDr _D^r. Take Θ¯=Θ¯r,K = ^r,K and Θi=Θir,K _i= _i^r,K, then ‖1S∑i∈Sr(PΘir,K(vi)−PΘ¯r,K(vi))‖2≤LΘ2(1S∑i∈Sr‖vi‖2)ΔDr.E \| 1S _i∈ S_r (P_ _i^r,K(v_i)-P_ ^r,K(v_i) ) \|^2≤ L_ ^2 ( 1S _i∈ S_rE\|v_i\|^2 ) _D^r. (19) B.4 A Decomposition of the Preconditioned Direction Fix the reference state Θ¯r,K ^r,K and decompose g~ir,k g_i^r,k as g~ir,k g_i^r,k =PΘ¯r,K(∇F(xr)) =P_ ^r,K\! (∇ F(x^r) ) +(PΘ¯r,K(∇Fi(xr))−PΘ¯r,K(∇F(xr)))⏟heterogeneity term + (P_ ^r,K\! (∇ F_i(x^r) )-P_ ^r,K\! (∇ F(x^r) ) )_ heterogeneity term +(PΘ¯r,K(gir,k)−PΘ¯r,K(∇Fi(xir,k)))⏟stochastic/local-drift term + (P_ ^r,K\! (g_i^r,k )-P_ ^r,K\! (∇ F_i(x_i^r,k) ) )_ stochastic/local-drift term +(PΘir,k(gir,k)−PΘ¯r,K(gir,k))⏟precond drift term. + (P_ _i^r,k\! (g_i^r,k )-P_ ^r,K\! (g_i^r,k ) )_ precond drift term. (20) Define the aggregated error Er:=1S∑i∈Sr∑k=0K−1[heteroir,k+stochir,k+pdriftir,k],Ur=KPΘ¯r,K(∇F(xr))+Er.E^r:= 1S _i∈ S_r _k=0^K-1 [ hetero_i^r,k+ stoch_i^r,k+ pdrift_i^r,k ], U^r=K\,P_ ^r,K\! (∇ F(x^r) )+E^r. (21) We first recall the definition of the server-aggregated direction: Ur:=1S∑i∈Sr∑k=0K−1g~ir,k.U^r\;:=\; 1S _i∈ S_r _k=0^K-1 g_i^r,k. (22) Fix a reference preconditioner state Θ¯r,K ^r,K. Using the decomposition in (21), i.e., g~ir,k=PΘ¯r,K(∇F(xr))+heteroir,k+stochir,k+pdriftir,k, g_i^r,k=P_ ^r,K\! (∇ F(x^r) )+ hetero_i^r,k+ stoch_i^r,k+ pdrift_i^r,k, (23) and substituting (23) into (22), we obtain Ur U^r =1S∑i∈Sr∑k=0K−1[PΘ¯r,K(∇F(xr))+heteroir,k+stochir,k+pdriftir,k] = 1S _i∈ S_r _k=0^K-1 [P_ ^r,K\! (∇ F(x^r) )+ hetero_i^r,k+ stoch_i^r,k+ pdrift_i^r,k ] =1S∑i∈Sr∑k=0K−1PΘ¯r,K(∇F(xr))⏟:=(⋆)+1S∑i∈Sr∑k=0K−1[heteroir,k+stochir,k+pdriftir,k]. = 1S _i∈ S_r _k=0^K-1P_ ^r,K\! (∇ F(x^r) )_:=\,( )\;+\; 1S _i∈ S_r _k=0^K-1 [ hetero_i^r,k+ stoch_i^r,k+ pdrift_i^r,k ]. (24) Since PΘ¯r,K(∇F(xr))P_ ^r,K\! (∇ F(x^r) ) does not depend on i nor k, we simplify (⋆)( ) as (⋆)=1S∑i∈Sr∑k=0K−1PΘ¯r,K(∇F(xr))=1S∑i∈Sr(∑k=0K−11)PΘ¯r,K(∇F(xr))=KPΘ¯r,K(∇F(xr)).( )= 1S _i∈ S_r _k=0^K-1P_ ^r,K\! (∇ F(x^r) )= 1S _i∈ S_r ( _k=0^K-11 )\,P_ ^r,K\! (∇ F(x^r) )=K\,P_ ^r,K\! (∇ F(x^r) ). (25) Define the aggregated error term Er:=1S∑i∈Sr∑k=0K−1[heteroir,k+stochir,k+pdriftir,k].E^r\;:=\; 1S _i∈ S_r _k=0^K-1 [ hetero_i^r,k+ stoch_i^r,k+ pdrift_i^r,k ]. (26) Substituting (25) and (26) into (24) yields Ur=KPΘ¯r,K(∇F(xr))+Er.U^r=K\,P_ ^r,K\! (∇ F(x^r) )+E^r. (27) B.5 One-Step Descent Inequality By L-smoothness of F, F(xr+1)≤F(xr)+⟨∇F(xr),xr+1−xr⟩+L2‖xr+1−xr‖2.F(x^r+1)≤ F(x^r)+ ∇ F(x^r),x^r+1-x^r + L2\|x^r+1-x^r\|^2. (28) Using (14), xr+1−xr=−ηUrx^r+1-x^r=-η U^r, F(xr+1)≤F(xr)−η⟨∇F(xr),Ur⟩+Lη22‖Ur‖2.F(x^r+1)≤ F(x^r)-η ∇ F(x^r),U^r + Lη^22\|U^r\|^2. (29) Plug (21) into the inner product: −η⟨∇F(xr),Ur⟩ -η ∇ F(x^r),U^r =−ηK⟨∇F(xr),PΘ¯r,K(∇F(xr))⟩−η⟨∇F(xr),Er⟩ =-η K ∇ F(x^r),P_ ^r,K\! (∇ F(x^r) ) -η ∇ F(x^r),E^r ≤−ηKμ‖∇F(xr)‖2−η⟨∇F(xr),Er⟩, ≤-η Kμ\|∇ F(x^r)\|^2-η ∇ F(x^r),E^r , (30) where we used (16). For the cross term, apply Young’s inequality with parameter μ/2μ/2: −η⟨∇F(xr),Er⟩≤η⋅μ4‖∇F(xr)‖2+η⋅1μ‖Er‖2.-η ∇ F(x^r),E^r ≤η· μ4\|∇ F(x^r)\|^2+η· 1μ\|E^r\|^2. (31) Combining (29)–(31), F(xr+1)≤F(xr)−ηK3μ4‖∇F(xr)‖2+η1μ‖Er‖2+Lη22‖Ur‖2.F(x^r+1)≤ F(x^r)-η K 3μ4\|∇ F(x^r)\|^2+η 1μ\|E^r\|^2+ Lη^22\|U^r\|^2. (32) B.6 Bounding ‖Er‖2E\|E^r\|^2 and ‖Ur‖2E\|U^r\|^2 We bound three components in ErE^r. (i) Heterogeneity term. By (16) boundedness and (A3), ‖PΘ¯r,K(∇Fi(xr))−PΘ¯r,K(∇F(xr))‖2≤M2‖∇Fi(xr)−∇F(xr)‖2≤M2σg2. \|P_ ^r,K\! (∇ F_i(x^r) )-P_ ^r,K\! (∇ F(x^r) ) \|^2≤ M^2\|∇ F_i(x^r)-∇ F(x^r)\|^2≤ M^2 _g^2. (33) (i) Stochastic/local-drift term. By (16) boundedness and (A2), ‖PΘ¯r,K(gir,k)−PΘ¯r,K(∇Fi(xir,k))‖2 \|P_ ^r,K(g_i^r,k)-P_ ^r,K(∇ F_i(x_i^r,k)) \|^2 ≤M2‖gir,k−∇Fi(xir,k)‖2 ≤ M^2\,E\|g_i^r,k-∇ F_i(x_i^r,k)\|^2 ≤M2σl2. ≤ M^2 _l^2. (34) (Any additional “local model drift” term caused by xir,k≠xrx_i^r,k≠ x^r can be absorbed into the constant by standard arguments; we keep the presentation clean by grouping it into the stochastic/local-drift bucket.) (i) Preconditioner drift term. Applying Lemma B.1 with vi=gir,kv_i=g_i^r,k and the instantiation (19), together with (A4), yields ‖1S∑i∈Sr(PΘir,k(gir,k)−PΘ¯r,K(gir,k))‖2≤LΘ2G2ΔDr.E \| 1S _i∈ S_r (P_ _i^r,k(g_i^r,k)-P_ ^r,K(g_i^r,k) ) \|^2≤ L_ ^2\,G^2\, _D^r. (35) Preconditioner drift term (detailed proof). Recall Lemma B.1: for any reference state Θ¯ , any local states Θii∈Sr\ _i\_i∈ S_r, and any vectors vii∈Sr\v_i\_i∈ S_r, ‖1S∑i∈Sr(PΘi(vi)−PΘ¯(vi))‖2≤LΘ2(1S∑i∈Sr‖vi‖2)(1S∑i∈Sr‖Θi−Θ¯‖2). \| 1S _i∈ S_r (P_ _i(v_i)-P_ (v_i) ) \|^2≤ L_ ^2 ( 1S _i∈ S_r\|v_i\|^2 ) ( 1S _i∈ S_r\| _i- \|^2 ). (36) We instantiate (36) by choosing Θi←Θir,k,Θ¯←Θ¯r,K,vi←gir,k. _i← _i^r,k, ← ^r,K, v_i← g_i^r,k. Then (36) gives ‖1S∑i∈Sr(PΘir,k(gir,k)−PΘ¯r,K(gir,k))‖2 \| 1S _i∈ S_r (P_ _i^r,k(g_i^r,k)-P_ ^r,K(g_i^r,k) ) \|^2 ≤LΘ2(1S∑i∈Sr‖gir,k‖2)(1S∑i∈Sr‖Θir,k−Θ¯r,K‖2). ≤ L_ ^2 ( 1S _i∈ S_r\|g_i^r,k\|^2 ) ( 1S _i∈ S_r\| _i^r,k- ^r,K\|^2 ). (37) Taking expectation on both sides and using linearity of expectation yields ‖1S∑i∈Sr(PΘir,k(gir,k)−PΘ¯r,K(gir,k))‖2 \| 1S _i∈ S_r (P_ _i^r,k(g_i^r,k)-P_ ^r,K(g_i^r,k) ) \|^2 ≤LΘ2[(1S∑i∈Sr‖gir,k‖2)(1S∑i∈Sr‖Θir,k−Θ¯r,K‖2)]. ≤ L_ ^2\,E [ ( 1S _i∈ S_r\|g_i^r,k\|^2 ) ( 1S _i∈ S_r\| _i^r,k- ^r,K\|^2 ) ]. (38) Next, we bound the first factor by Assumption (A4): ‖gir,k‖2≤G2E\|g_i^r,k\|^2≤ G^2 for all i,r,ki,r,k. Hence, (1S∑i∈Sr‖gir,k‖2)=1S∑i∈Sr‖gir,k‖2≤G2.E ( 1S _i∈ S_r\|g_i^r,k\|^2 )= 1S _i∈ S_rE\|g_i^r,k\|^2≤ G^2. (39) Finally, to express the second factor in terms of the drift metric, we use the instantiation (19) with Θ¯=Θ¯r,K = ^r,K. If the drift metric is defined at step k as ΔDr,k:=1S∑i∈Sr‖Θir,k−Θ¯r,K‖2, _D^r,k:= 1S _i∈ S_rE \| _i^r,k- ^r,K \|^2, (40) then (38)–(39) imply ‖1S∑i∈Sr(PΘir,k(gir,k)−PΘ¯r,K(gir,k))‖2≤LΘ2G2ΔDr,k.E \| 1S _i∈ S_r (P_ _i^r,k(g_i^r,k)-P_ ^r,K(g_i^r,k) ) \|^2≤ L_ ^2\,G^2\, _D^r,k. (41) In particular, if we upper bound ΔDr,k≤ΔDr _D^r,k≤ _D^r for all k∈0,…,K−1k∈\0,…,K-1\ (e.g., by monotonicity or by defining ΔDr:=max0≤k≤K⁡ΔDr,k _D^r:= _0≤ k≤ K _D^r,k), we obtain the stated bound: ‖1S∑i∈Sr(PΘir,k(gir,k)−PΘ¯r,K(gir,k))‖2≤LΘ2G2ΔDr.E \| 1S _i∈ S_r (P_ _i^r,k(g_i^r,k)-P_ ^r,K(g_i^r,k) ) \|^2≤ L_ ^2\,G^2\, _D^r. (42) Remark (optional). If one prefers to keep the dependence on k, then (41) is the tight form, and the analysis can carry ΔDr,k _D^r,k through the subsequent steps. Aggregate bound. Using ‖∑t=1Tat‖2≤T∑t=1T‖at‖2\| _t=1^Ta_t\|^2≤ T _t=1^T\|a_t\|^2 and the fact that averaging over S clients reduces variance by S, we obtain (for some universal constant c>0c>0) ‖Er‖2≤cK(M2σl2S+M2σg2S+LΘ2G2ΔDr).E\|E^r\|^2≤ c\,K ( M^2 _l^2S+ M^2 _g^2S+L_ ^2G^2 _D^r ). (43) Moreover, by ‖a+b‖2≤2‖a‖2+2‖b‖2\|a+b\|^2≤ 2\|a\|^2+2\|b\|^2 and (16), ‖Ur‖2 \|U^r\|^2 =‖KPΘ¯r,K(∇F(xr))+Er‖2 =E \|KP_ ^r,K(∇ F(x^r))+E^r \|^2 ≤2K2‖PΘ¯r,K(∇F(xr))‖2+2‖Er‖2 ≤ 2K^2E\|P_ ^r,K(∇ F(x^r))\|^2+2E\|E^r\|^2 ≤2K2M2‖∇F(xr)‖2+2‖Er‖2. ≤ 2K^2M^2E\|∇ F(x^r)\|^2+2E\|E^r\|^2. (44) B.7 Per-Round Recursion with an Explicit Drift Penalty Taking expectation of (32) and plugging (43)–(44), we obtain F(xr+1) (x^r+1) ≤F(xr)−ηK(3μ4−LηKM2)‖∇F(xr)‖2 (x^r)-η K ( 3μ4-Lη KM^2 )E\|∇ F(x^r)\|^2 +cη(1μ+Lη)K(M2(σl2+σg2)S+LΘ2G2ΔDr). +c\,η ( 1μ+Lη )K ( M^2( _l^2+ _g^2)S+L_ ^2G^2 _D^r ). (45) Choose η such that LηKM2≤μ4Lη KM^2≤ μ4, i.e., η≤μ4LKM2.η≤ μ4LKM^2. (46) Then the descent coefficient is positive and (45) simplifies to F(xr+1)≤F(xr)−μ2ηK‖∇F(xr)‖2+C0Lη2K(σl2+σg2S+LΘ2G2M2⏟:=κΘΔDr),EF(x^r+1) (x^r)- μ2η K\,E\|∇ F(x^r)\|^2+C_0\,Lη^2K ( _l^2+ _g^2S+ L_ ^2G^2M^2_:=\; _ _D^r ), (47) where C0>0C_0>0 is an absolute constant (absorbing M2M^2 and μ). B.8 Telescoping and Final Rate (FedSOA with Drift) Sum (47) for r=0,…,R−1r=0,…,R-1 and use Δ:=F(x0)−F⋆ :=F(x^0)-F : 1R∑r=0R−1‖∇F(xr)‖2 1R _r=0^R-1E\|∇ F(x^r)\|^2 ≤2ΔμηKR+C1Lημ(σl2+σg2S+κΘΔ¯D), ≤ 2 μη KR+C_1\, Lημ ( _l^2+ _g^2S+ _ \, _D ), (48) where Δ¯D:=1R∑r=0R−1ΔDr _D:= 1R _r=0^R-1 _D^r and C1C_1 is an absolute constant. Optimizing the RHS over η (subject to (46)) yields the standard non-convex form 1R∑r=0R−1‖∇F(xr)‖2≲(LΔR+LΔR⋅σl2+σg2+SκΘΔ¯DSK), 1R _r=0^R-1E\|∇ F(x^r)\|^2\; \;O\! ( L R+ L R· _l^2+ _g^2+S _ \, _DSK ), (49) where we use ≲ to hide absolute constants depending only on (μ,M)(μ,M). Interpretation. Eq. 49 contains an explicit drift penalty: the larger Δ¯D _D is, the larger the “effective noise” is, hence the slower the convergence. Appendix C Proof of Theorem 5.7 (FedPAC) C.1 Further Alignment to Theorem 5.7 (FedPAC): Why σg2 _g^2 Disappears FedPAC introduces two mechanisms: (i) Preconditioner Alignment. At the beginning of round r, each participating client warm-starts from a shared reference Θir,0←Θr _i^r,0← ^r, and the server aggregates states after local steps. This reduces preconditioner mismatch and thus controls ΔDr _D^r. (i) Local Preconditioner Correction. Each local step uses a convex combination of local preconditioned direction and an estimated global direction: xir,k+1=xir,k−η[(1−β)g~ir,k+βgGr],β∈[0,1],x_i^r,k+1=x_i^r,k-η [(1-β) g_i^r,k+β g_G^r ], β∈[0,1], (50) where gGrg_G^r is the estimated global update from the previous round, e.g. gGr:=−1SKη∑i∈Sr−1(xir−1,K−xir−1,0).g_G^r:=- 1SKη _i∈ S_r-1 (x_i^r-1,K-x_i^r-1,0 ). (51) Algorithm 3 Federated Preconditioner Alignment and Correction Framework (FedPAC) 0: Per client, we maintain: ir,k _i^r,k in local. Hyperparameters: learning rate η, communication rounds R, local updates K, the number of client N. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈ri _r in parallel do 3: ir,0←r; _i^r,0← ^r; 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bir,kB_i^r,k; 6: ir,k∈ℝm×n←∇Fi(ir,k;ξir,k) g_i^r,k ^m× n←∇ F_i( x_i^r,k; _i^r,k); 7: ir,k+1=UpdateState(ir,k,ir,k); _i^r,k+1=UpdateState ( _i^r,k, g_i^r,k ); 8: ~ir,k←ir,k(ir,k); g_i^r,k _ _i^r,k ( g_i^r,k ); 9: ir,k+1=ir,k−ηl[(1−β)~ir,k+βGr] x^r,k+1_i\!=\! x^r,k_i\!\!-\! _l[(1\!-\!β) g_i^r,k\!+\!β g_G^r]; 10: end for 11: Δir≔ir,K−r x_i^r x_i^r,K- x^r; 12: Client i communicate (ir,ir,K)( x_i^r, _i^r,K) to Server; 13: end for 14: Gr+1=−1SKη∑i=1S(ir,K−ir,0) g_G^r+1=- 1SKη _i=1^S( x^r,K_i- x^r,0_i); 15: r+1=r−γGr+1; x^r+1= x^r-γ g_G^r+1; 16:r+1=1|r|∑i∈rir,K; ^r+1\;=\; 1|S_r| _i _r _i^r,K; 17: Server broadcasts (r+1,r+1,Gr+1);( x^r+1, _r+1, g_G^r+1); 18: end for C.2 Further Alignment to Theorem 5.5 (FedPAC): Why σg2 _g^2 Disappears (Detailed Proof) Step 0: Why σg2 _g^2 appears in FedSOA. In the FedSOA proof, we explicitly insert and subtract ∇Fi(xr)∇ F_i(x^r), producing the term PΘ¯r,K(∇Fi(xr))−PΘ¯r,K(∇F(xr))P_ ^r,K(∇ F_i(x^r))-P_ ^r,K(∇ F(x^r)). Bounding its second moment requires Assumption (A3), i.e., ‖∇Fi(x)−∇F(x)‖2≤σg2\|∇ F_i(x)-∇ F(x)\|^2≤ _g^2, which yields the σg2 _g^2 term. Step 1: Use a global-centered variance assumption instead of (A3). For FedPAC, we do not introduce ∇Fi(xr)∇ F_i(x^r) in the decomposition. Instead, we impose the following global-centered noise condition (which does not require bounded heterogeneity): C.3 Recursive Analysis to Eliminate Data Heterogeneity (Final) Goal. We present an appendix-ready proof template showing how to eliminate the explicit heterogeneity term (i.e., no σg2 _g^2) via a recursive global-direction estimator, in the same spirit as DP-FedPGN-style recursion. The remaining terms depend on σl2 _l^2 (variance reduced by SKSK) and the preconditioner drift penalty (controlled by ΔDr _D^r). Notation. Let f(x):=1N∑i=1Nfi(x)f(x):= 1N _i=1^Nf_i(x) be L-smooth. In round r, a set SrS_r of S clients participates and each performs K local steps. Let Θ¯r,K ^r,K be the reference preconditioner state in round r (e.g., the aggregated state), and define the local preconditioned stochastic gradient vir,k:=PΘir,k(gir,k),gir,k:=∇fi(xir,k;ξir,k).v_i^r,k:=P_ _i^r,k\! (g_i^r,k ), g_i^r,k:=∇ f_i(x_i^r,k; _i^r,k). Define the per-round observation (average over clients and local steps) v^r:=1SK∑i∈Sr∑k=0K−1vir,k. v^r\;:=\; 1SK _i∈ S_r _k=0^K-1v_i^r,k. (52) Recursive global-direction estimator. We maintain a global direction estimator grr≥0\g^r\_r≥ 0: gr+1:=(1−β)gr+βv^r,β∈(0,1],g^r+1\;:=\;(1-β)g^r+β v^r, β∈(0,1], (53) and update the global model by xr+1=xr−γgr+1.x^r+1\;=\;x^r-γ g^r+1. (54) Assumptions. We use: 1. (Smoothness) f is L-smooth. 2. (Local stochastic variance) for all i,r,ki,r,k, ‖gir,k−∇fi(xir,k)‖2≤σl2E\|g_i^r,k-∇ f_i(x_i^r,k)\|^2≤ _l^2. 3. (Preconditioner boundedness) for all Θ,v ,v, ‖PΘ(v)‖≤M‖v‖\|P_ (v)\|≤ M\|v\|. 4. (Preconditioner drift Lipschitz) for all Θ,Θ′,v , ,v, ‖PΘ(v)−PΘ′(v)‖≤LΘ‖Θ−Θ′‖⋅‖v‖\|P_ (v)-P_ (v)\|≤ L_ \| - \|·\|v\|. 5. (Second-moment bound) supi,r,k‖gir,k‖2≤G2 _i,r,kE\|g_i^r,k\|^2≤ G^2. Importantly, we do not assume bounded heterogeneity ‖∇fi(x)−∇f(x)‖2≤σg2\|∇ f_i(x)-∇ f(x)\|^2≤ _g^2. Local trajectory drift measure. Define Ur:=1SK∑i∈Sr∑k=0K−1‖xir,k−xr‖2.U_r:= 1SK _i∈ S_r _k=0^K-1E\|x_i^r,k-x^r\|^2. (55) Preconditioner drift metric. Define the per-round drift ΔDr:=1S∑i∈Sr‖Θir,K−Θ¯r,K‖2. _D^r:= 1S _i∈ S_rE\| _i^r,K- ^r,K\|^2. (56) Lemma C.1 (Drift-induced preconditioned disagreement). For any vectors uii∈Sr\u_i\_i∈ S_r, ‖1S∑i∈Sr(PΘir,k(ui)−PΘ¯r,K(ui))‖2≤LΘ2(1S∑i∈Sr‖ui‖2)ΔDr.E \| 1S _i∈ S_r (P_ _i^r,k(u_i)-P_ ^r,K(u_i) ) \|^2\;≤\;L_ ^2 ( 1S _i∈ S_rE\|u_i\|^2 ) _D^r. (57) In particular, using ‖ui‖2≤G2E\|u_i\|^2≤ G^2 gives ‖1S∑i∈Sr(PΘir,k(ui)−PΘ¯r,K(ui))‖2≤LΘ2G2ΔDr.E \| 1S _i∈ S_r (P_ _i^r,k(u_i)-P_ ^r,K(u_i) ) \|^2\;≤\;L_ ^2G^2\, _D^r. (58) Lemma C.2 (Descent with inexact direction). Under L-smoothness, if γL≤124γ L≤ 124, then for all r≥0r≥ 0, [f(xr+1)]≤[f(xr)]−11γ24‖∇f(xr)‖2+13γ24ℰ~r,E [f(x^r+1) ] [f(x^r) ]- 11γ24\,E\|∇ f(x^r)\|^2+ 13γ24\, E_r, (59) where ℰ~r:=‖∇f(xr)−gr+1‖2 E_r:=E\|∇ f(x^r)-g^r+1\|^2. Proof. Since f is L-smooth, for any x,yx,y we have f(y)≤f(x)+⟨∇f(x),y−x⟩+L2‖y−x‖2.f(y)≤ f(x)+ ∇ f(x),y-x + L2\|y-x\|^2. Apply it with x=xrx=x^r and y=xr+1=xr−γgr+1y=x^r+1=x^r-γ g^r+1: f(xr+1) f(x^r+1) ≤f(xr)+⟨∇f(xr),−γgr+1⟩+L2γ2‖gr+1‖2 ≤ f(x^r)+ ∇ f(x^r),-γ g^r+1 + L2γ^2\|g^r+1\|^2 =f(xr)−γ‖∇f(xr)‖2+γ⟨∇f(xr),∇f(xr)−gr+1⟩+L2γ2‖gr+1‖2. =f(x^r)-γ\|∇ f(x^r)\|^2+γ ∇ f(x^r),∇ f(x^r)-g^r+1 + L2γ^2\|g^r+1\|^2. (60) For the cross term, by Young’s inequality ⟨a,b⟩≤12‖a‖2+12‖b‖2 a,b ≤ 12\|a\|^2+ 12\|b\|^2, ⟨∇f(xr),∇f(xr)−gr+1⟩≤12‖∇f(xr)‖2+12‖∇f(xr)−gr+1‖2. ∇ f(x^r),∇ f(x^r)-g^r+1 ≤ 12\|∇ f(x^r)\|^2+ 12\|∇ f(x^r)-g^r+1\|^2. (61) For the last term, use ‖gr+1‖2=‖∇f(xr)−(∇f(xr)−gr+1)‖2≤2‖∇f(xr)‖2+2‖∇f(xr)−gr+1‖2\|g^r+1\|^2=\|∇ f(x^r)-(∇ f(x^r)-g^r+1)\|^2≤ 2\|∇ f(x^r)\|^2+2\|∇ f(x^r)-g^r+1\|^2 to get L2γ2‖gr+1‖2≤Lγ2‖∇f(xr)‖2+Lγ2‖∇f(xr)−gr+1‖2. L2γ^2\|g^r+1\|^2≤ Lγ^2\|∇ f(x^r)\|^2+Lγ^2\|∇ f(x^r)-g^r+1\|^2. (62) Substituting (61) and (62) into (60) yields f(xr+1) f(x^r+1) ≤f(xr)−γ‖∇f(xr)‖2+γ2‖∇f(xr)‖2+γ2‖∇f(xr)−gr+1‖2 ≤ f(x^r)-γ\|∇ f(x^r)\|^2+ γ2\|∇ f(x^r)\|^2+ γ2\|∇ f(x^r)-g^r+1\|^2 +Lγ2‖∇f(xr)‖2+Lγ2‖∇f(xr)−gr+1‖2 +Lγ^2\|∇ f(x^r)\|^2+Lγ^2\|∇ f(x^r)-g^r+1\|^2 =f(xr)−(γ2−Lγ2)‖∇f(xr)‖2+(γ2+Lγ2)‖∇f(xr)−gr+1‖2. =f(x^r)- ( γ2-Lγ^2 )\|∇ f(x^r)\|^2+ ( γ2+Lγ^2 )\|∇ f(x^r)-g^r+1\|^2. (63) If γL≤124γ L≤ 124, then γ2−Lγ2=γ(12−γL)≥γ(12−124)=11γ24,γ2+Lγ2=γ(12+γL)≤γ(12+124)=13γ24. γ2-Lγ^2=γ ( 12-γ L )≥γ ( 12- 124 )= 11γ24, γ2+Lγ^2=γ ( 12+γ L )≤γ ( 12+ 124 )= 13γ24. Thus (63) implies f(xr+1)≤f(xr)−11γ24‖∇f(xr)‖2+13γ24‖∇f(xr)−gr+1‖2.f(x^r+1)≤ f(x^r)- 11γ24\|∇ f(x^r)\|^2+ 13γ24\|∇ f(x^r)-g^r+1\|^2. Taking expectation on both sides and recalling ℰ~r:=‖∇f(xr)−gr+1‖2 E_r:=E\|∇ f(x^r)-g^r+1\|^2 completes the proof. ∎ Lemma C.3 (Recursive estimator error (no explicit σg2 _g^2)). Let ℰ~r:=‖∇f(xr)−gr+1‖2 E_r:=E\|∇ f(x^r)-g^r+1\|^2 with gr+1g^r+1 defined in (53)–(52). If γL≤β6γ L≤ β6, then for all r≥1r≥ 1, ℰ~r≤(1−8β9)ℰ~r−1+C1γ2L2β‖∇f(xr−1)‖2+C2β2σl2SK+C3βL2Ur+C4βLΘ2G2ΔDr, E_r\;≤\; (1- 8β9 ) E_r-1+ C_1γ^2L^2β\,E\|∇ f(x^r-1)\|^2+ C_2β^2 _l^2SK+C_3β L^2U_r+C_4β\,L_ ^2G^2\, _D^r, (64) and for r=0r=0, ℰ~0≤(1−β)ℰ~−1+C2β2σl2SK+C3βL2U0+C4βLΘ2G2ΔD0, E_0\;≤\;(1-β) E_-1+ C_2β^2 _l^2SK+C_3β L^2U_0+C_4β\,L_ ^2G^2\, _D^0, (65) where C1,C2,C3,C4>0C_1,C_2,C_3,C_4>0 are absolute constants depending only on (M)(M). Proof. We follow a DP-FedPGN-style recursion and avoid introducing ∇fi−∇f∇ f_i-∇ f explicitly. Step 1: Expand the recursion. Define δr:=∇f(xr)−gr+1 _r:=∇ f(x^r)-g^r+1 so that ℰ~r=‖δr‖2 E_r=E\| _r\|^2. From (53) and (52), δr _r =∇f(xr)−(1−β)gr−βv^r =∇ f(x^r)-(1-β)g^r-β v^r =(1−β)(∇f(xr)−gr)+β(∇f(xr)−v^r). =(1-β) (∇ f(x^r)-g^r )+β (∇ f(x^r)- v^r ). (66) Add and subtract ∇f(xr−1)∇ f(x^r-1) inside the first bracket: ∇f(xr)−gr ∇ f(x^r)-g^r =∇f(xr)−∇f(xr−1)⏟Ar+∇f(xr−1)−gr⏟=δr−1. = ∇ f(x^r)-∇ f(x^r-1)_A_r+ ∇ f(x^r-1)-g^r_= _r-1. (67) Plugging (67) into (66) gives δr=(1−β)δr−1+(1−β)Ar+βBr,Br:=∇f(xr)−v^r. _r=(1-β) _r-1+(1-β)A_r+β B_r, B_r:=∇ f(x^r)- v^r. (68) Step 2: Square and apply the same AM-GM pattern. Using ‖a+b‖2≤(1+β2)‖a‖2+(1+2β)‖b‖2\|a+b\|^2≤(1+ β2)\|a\|^2+(1+ 2β)\|b\|^2 twice (as in your proof), one obtains ‖δr‖2≤(1+β2)(1−β)2‖δr−1‖2+Cβ‖Ar‖2+Cβ2‖Br‖2,E\| _r\|^2≤ (1+ β2 )(1-β)^2E\| _r-1\|^2+ Cβ\,E\|A_r\|^2+Cβ^2\,E\|B_r\|^2, (69) for an absolute constant C>0C>0. Step 3: Bound ‖Ar‖2E\|A_r\|^2 by smoothness and the update. By L-smoothness, ‖∇f(xr)−∇f(xr−1)‖≤L‖xr−xr−1‖\|∇ f(x^r)-∇ f(x^r-1)\|≤ L\|x^r-x^r-1\| and xr−xr−1=−γgrx^r-x^r-1=-γ g^r, hence ‖Ar‖2 \|A_r\|^2 ≤L2γ2‖gr‖2≤2L2γ2(‖∇f(xr−1)‖2+‖∇f(xr−1)−gr‖2)=2L2γ2(‖∇f(xr−1)‖2+ℰ~r−1). ≤ L^2γ^2\,E\|g^r\|^2≤ 2L^2γ^2 (E\|∇ f(x^r-1)\|^2+E\|∇ f(x^r-1)-g^r\|^2 )=2L^2γ^2 (E\|∇ f(x^r-1)\|^2+ E_r-1 ). (70) Step 4: Bound ‖Br‖2E\|B_r\|^2 without σg2 _g^2. Recall Br=∇f(xr)−v^rB_r=∇ f(x^r)- v^r and v^r=1SK∑i,kPΘir,k(gir,k) v^r= 1SK _i,kP_ _i^r,k(g_i^r,k). Add and subtract the reference-preconditioned terms: Br B_r =∇f(xr)−1SK∑i,kPΘ¯r,K(gir,k)⏟Br(1)+1SK∑i,k(PΘ¯r,K(gir,k)−PΘir,k(gir,k))⏟Br(2). = ∇ f(x^r)- 1SK _i,kP_ ^r,K\! (g_i^r,k )_B_r^(1)+ 1SK _i,k (P_ ^r,K(g_i^r,k)-P_ _i^r,k(g_i^r,k) )_B_r^(2). (71) (i) Stochastic + trajectory part Br(1)B_r^(1). Insert and subtract ∇f(xir,k)∇ f(x_i^r,k) and use ‖PΘ¯(v)‖≤M‖v‖\|P_ (v)\|≤ M\|v\|: ‖Br(1)‖2 \|B_r^(1)\|^2 =‖1SK∑i,k(∇f(xr)−PΘ¯r,K(gir,k))‖2 = \| 1SK _i,k (∇ f(x^r)-P_ ^r,K(g_i^r,k) ) \|^2 ≤2SK∑i,k‖∇f(xr)−PΘ¯r,K(∇f(xir,k))‖2+2SK∑i,k‖PΘ¯r,K(∇f(xir,k)−gir,k)‖2 ≤ 2SK _i,k \|∇ f(x^r)-P_ ^r,K(∇ f(x_i^r,k)) \|^2+ 2SK _i,k \|P_ ^r,K(∇ f(x_i^r,k)-g_i^r,k) \|^2 ≤2M2SK∑i,k‖∇f(xr)−∇f(xir,k)‖2+2M2SK∑i,k‖∇f(xir,k)−gir,k‖2. ≤ 2M^2SK _i,k\|∇ f(x^r)-∇ f(x_i^r,k)\|^2+ 2M^2SK _i,k\|∇ f(x_i^r,k)-g_i^r,k\|^2. (72) Taking expectation, using local variance ‖gir,k−∇fi(xir,k)‖2≤σl2E\|g_i^r,k-∇ f_i(x_i^r,k)\|^2≤ _l^2 and L-smoothness ‖∇f(xir,k)−∇f(xr)‖2≤L2‖xir,k−xr‖2\|∇ f(x_i^r,k)-∇ f(x^r)\|^2≤ L^2\|x_i^r,k-x^r\|^2, and absorbing the client/gradient mismatch into the recursion (rather than a σg2 _g^2 constant), we get the standard bound ‖Br(1)‖2≤C(σl2SK+L2Ur),E\|B_r^(1)\|^2≤ C ( _l^2SK+L^2U_r ), (73) where UrU_r is defined in (55) and C>0C>0 depends only on M. (i) Preconditioner drift part Br(2)B_r^(2). By Lemma C.1 with ui=gir,ku_i=g_i^r,k and (58), ‖Br(2)‖2≤CLΘ2G2ΔDr.E\|B_r^(2)\|^2≤ C\,L_ ^2G^2\, _D^r. (74) Combining (71)–(74) gives ‖Br‖2≤C(σl2SK+L2Ur+LΘ2G2ΔDr).E\|B_r\|^2≤ C ( _l^2SK+L^2U_r+L_ ^2G^2\, _D^r ). (75) Step 5: Conclude the recursion and absorb ℰ~r−1 E_r-1. Plug (70) and (75) into (69). Using γL≤β/6γ L≤β/6 to absorb the ℰ~r−1 E_r-1 term into the contraction, we obtain (64)–(65) with absolute constants. ∎ Theorem C.4 (Non-convex convergence without explicit heterogeneity). Under the assumptions above, take g0=0g^0=0 and choose β≍min⁡1,SKLΔσl2R,γ=min⁡124L,β6L,β \1, SKL _l^2R \, γ= \ 124L, β6L \, and ηKLη KL sufficiently small so that the local drift bound UrU_r is controlled (as in the auxiliary lemma). Then we have 1R∑r=0R−1‖∇f(xr)‖2≲LΔR+LΔσl2SKR+γLΘ2G2⋅1R∑r=0R−1ΔDr⏟drift penalty. 1R _r=0^R-1E\|∇ f(x^r)\|^2\; \; L R+ L \, _l^2SKR\;+\; γ L_ ^2G^2· 1R _r=0^R-1 _D^r_drift penalty. (76) In particular, under alignment ΔDr≈0 _D^r≈ 0, the last term vanishes and the rate depends only on σl2 _l^2. Proof. Summing the descent inequality in Lemma C.2 over r=0,…,R−1r=0,…,R-1 yields 11γ24∑r=0R−1‖∇f(xr)‖2≤Δ+13γ24∑r=0R−1ℰ~r,Δ:=f(x0)−f⋆. 11γ24 _r=0^R-1E\|∇ f(x^r)\|^2≤ + 13γ24 _r=0^R-1 E_r, :=f(x^0)-f . (77) It remains to bound ∑r=0R−1ℰ~r _r=0^R-1 E_r. Step 1: Sum the recursion for ℰ~r E_r. From Lemma C.3, for r≥1r≥ 1, ℰ~r≤(1−8β9)ℰ~r−1+C1γ2L2β‖∇f(xr−1)‖2+C2β2σl2SK+C3βL2Ur+C4βLΘ2G2ΔDr. E_r≤ (1- 8β9 ) E_r-1+ C_1γ^2L^2β\,E\|∇ f(x^r-1)\|^2+ C_2β^2 _l^2SK+C_3β L^2U_r+C_4β\,L_ ^2G^2\, _D^r. Summing both sides over r=1,…,R−1r=1,…,R-1 gives ∑r=1R−1ℰ~r _r=1^R-1 E_r ≤(1−8β9)∑r=1R−1ℰ~r−1+C1γ2L2β∑r=1R−1‖∇f(xr−1)‖2+C2β2σl2SK(R−1) ≤ (1- 8β9 ) _r=1^R-1 E_r-1+ C_1γ^2L^2β _r=1^R-1E\|∇ f(x^r-1)\|^2+ C_2β^2 _l^2SK(R-1) +C3βL2∑r=1R−1Ur+C4βLΘ2G2∑r=1R−1ΔDr. +C_3β L^2 _r=1^R-1U_r+C_4β L_ ^2G^2 _r=1^R-1 _D^r. (78) Re-index ∑r=1R−1ℰ~r−1=∑r=0R−2ℰ~r _r=1^R-1 E_r-1= _r=0^R-2 E_r and ∑r=1R−1‖∇f(xr−1)‖2=∑r=0R−2‖∇f(xr)‖2 _r=1^R-1E\|∇ f(x^r-1)\|^2= _r=0^R-2E\|∇ f(x^r)\|^2. Also note ∑r=1R−1ℰ~r≤∑r=0R−1ℰ~r _r=1^R-1 E_r≤ _r=0^R-1 E_r and ∑r=0R−2ℰ~r≤∑r=0R−1ℰ~r _r=0^R-2 E_r≤ _r=0^R-1 E_r. Thus (78) implies ∑r=0R−1ℰ~r _r=0^R-1 E_r ≤(1−8β9)∑r=0R−1ℰ~r+ℰ~0+C1γ2L2β∑r=0R−1‖∇f(xr)‖2+C2β2σl2SKR ≤ (1- 8β9 ) _r=0^R-1 E_r+ E_0+ C_1γ^2L^2β _r=0^R-1E\|∇ f(x^r)\|^2+ C_2β^2 _l^2SKR +C3βL2∑r=0R−1Ur+C4βLΘ2G2∑r=0R−1ΔDr. +C_3β L^2 _r=0^R-1U_r+C_4β L_ ^2G^2 _r=0^R-1 _D^r. (79) Move the contraction term to the left: 8β9∑r=0R−1ℰ~r≤ℰ~0+C1γ2L2β∑r=0R−1‖∇f(xr)‖2+C2β2σl2SKR+C3βL2∑r=0R−1Ur+C4βLΘ2G2∑r=0R−1ΔDr. 8β9 _r=0^R-1 E_r≤ E_0+ C_1γ^2L^2β _r=0^R-1E\|∇ f(x^r)\|^2+ C_2β^2 _l^2SKR+C_3β L^2 _r=0^R-1U_r+C_4β L_ ^2G^2 _r=0^R-1 _D^r. (80) Therefore, ∑r=0R−1ℰ~r≤98βℰ~0+9C18γ2L2β2∑r=0R−1‖∇f(xr)‖2+9C28βσl2SKR+9C38L2∑r=0R−1Ur+9C48LΘ2G2∑r=0R−1ΔDr. _r=0^R-1 E_r≤ 98β E_0+ 9C_18 γ^2L^2β^2 _r=0^R-1E\|∇ f(x^r)\|^2+ 9C_28 β _l^2SKR+ 9C_38L^2 _r=0^R-1U_r+ 9C_48L_ ^2G^2 _r=0^R-1 _D^r. (81) Step 2: Control ℰ~0 E_0 by ℰ~−1 E_-1. From Lemma C.3 at r=0r=0, ℰ~0≤(1−β)ℰ~−1+C2β2σl2SK+C3βL2U0+C4βLΘ2G2ΔD0≤ℰ~−1+C2β2σl2SK+C3βL2U0+C4βLΘ2G2ΔD0. E_0≤(1-β) E_-1+ C_2β^2 _l^2SK+C_3β L^2U_0+C_4β L_ ^2G^2 _D^0≤ E_-1+ C_2β^2 _l^2SK+C_3β L^2U_0+C_4β L_ ^2G^2 _D^0. If g0=0g^0=0, then ℰ~−1=‖∇f(x0)−g0‖2=‖∇f(x0)‖2≤2L(f(x0)−f⋆)=2LΔ E_-1=E\|∇ f(x^0)-g^0\|^2=\|∇ f(x^0)\|^2≤ 2L (f(x^0)-f )=2L (by standard smoothness inequality). Hence ℰ~0≤ 2LΔ+C2β2σl2SK+C3βL2U0+C4βLΘ2G2ΔD0. E_0\;≤\;2L + C_2β^2 _l^2SK+C_3β L^2U_0+C_4β L_ ^2G^2 _D^0. (82) Step 3: Plug (81) into (77) and isolate ∑‖∇f(xr)‖2Σ\|∇ f(x^r)\|^2. Substitute (81) into (77): 11γ24∑r=0R−1‖∇f(xr)‖2 11γ24 _r=0^R-1E\|∇ f(x^r)\|^2 ≤Δ+13γ24[98βℰ~0+9C18γ2L2β2∑r=0R−1∥∇f(xr)∥2+9C28βσl2SKR ≤ + 13γ24 [ 98β E_0+ 9C_18 γ^2L^2β^2 _r=0^R-1E\|∇ f(x^r)\|^2+ 9C_28 β _l^2SKR +9C38L2∑r=0R−1Ur+9C48LΘ2G2∑r=0R−1ΔDr]. + 9C_38L^2 _r=0^R-1U_r+ 9C_48L_ ^2G^2 _r=0^R-1 _D^r ]. (83) Move the term containing ∑‖∇f(xr)‖2 \|∇ f(x^r)\|^2 on the RHS to the LHS. Specifically, assume γL≤β/6γ L≤β/6 so that 13γ24⋅9C18⋅γ2L2β2≤12⋅11γ24, 13γ24· 9C_18· γ^2L^2β^2\;≤\; 12· 11γ24, (84) which can always be ensured by taking γL≲βγ L β (as in the theorem statement). Then (83) implies 11γ48∑r=0R−1‖∇f(xr)‖2 11γ48 _r=0^R-1E\|∇ f(x^r)\|^2 ≤Δ+13γ24⋅98βℰ~0+13γ24⋅9C28βσl2SKR ≤ + 13γ24· 98β E_0+ 13γ24· 9C_28 β _l^2SKR +13γ24⋅9C38L2∑r=0R−1Ur+13γ24⋅9C48LΘ2G2∑r=0R−1ΔDr. + 13γ24· 9C_38L^2 _r=0^R-1U_r+ 13γ24· 9C_48L_ ^2G^2 _r=0^R-1 _D^r. (85) Divide both sides by R and by γ: 1R∑r=0R−1‖∇f(xr)‖2 1R _r=0^R-1E\|∇ f(x^r)\|^2 ≤C(ΔγR+ℰ~0βR+βσl2SK+L2R∑r=0R−1Ur+LΘ2G2R∑r=0R−1ΔDr), ≤ C ( γ R+ E_0β R+ β _l^2SK+ L^2R _r=0^R-1U_r+ L_ ^2G^2R _r=0^R-1 _D^r ), (86) where C>0C>0 is an absolute constant. Step 4: Remove ℰ~0 E_0 and handle UrU_r (auxiliary bound). Plug (82) into (86): 1R∑r=0R−1‖∇f(xr)‖2 1R _r=0^R-1E\|∇ f(x^r)\|^2 ≤C(ΔγR+LΔβR+βσl2SK+L2R∑r=0R−1Ur+LΘ2G2R∑r=0R−1ΔDr)+C⋅βσl2SK⋅1R, ≤ C ( γ R+ L β R+ β _l^2SK+ L^2R _r=0^R-1U_r+ L_ ^2G^2R _r=0^R-1 _D^r )+C· β _l^2SK· 1R, (87) where the last tiny term can be absorbed into βσl2/(SK)β _l^2/(SK). Next, apply the auxiliary bound on UrU_r (as in your reference proof) to ensure L2R∑r=0R−1Ur≤C′βσl2SK, L^2R _r=0^R-1U_r\;≤\;C β _l^2SK, (88) by choosing ηKLη KL sufficiently small. This yields 1R∑r=0R−1‖∇f(xr)‖2≤C(ΔγR+LΔβR+βσl2SK+LΘ2G2R∑r=0R−1ΔDr). 1R _r=0^R-1E\|∇ f(x^r)\|^2\;≤\;C ( γ R+ L β R+ β _l^2SK+ L_ ^2G^2R _r=0^R-1 _D^r ). (89) Step 5: Choose (β,γ)(β,γ) and simplify to the final rate. Choose γ=min⁡124L,β6Lγ= \ 124L, β6L\ and β≍min⁡1,SKLΔσl2R.β \1, SKL _l^2R \. Then Δ/(γR)≲LΔ/R /(γ R) L /R and LΔβR+βσl2SK≲LΔσl2SKR. L β R+ β _l^2SK\; \; L _l^2SKR. Plugging these into (89) yields 1R∑r=0R−1‖∇f(xr)‖2≲LΔR+LΔσl2+κΘ⋅Δ¯DSKR, 1R _r=0^R-1E\|∇ f(x^r)\|^2\; \; L R+ L \, _l^2+ _ · _DSKR, which is exactly (76). ∎ Appendix D Appendix B: Experimental Setup D.1 Setting for ResNet-18 Table 7: A detailed summary of 100 and Tiny-ImageNet: number of classes, image size, and dataset splits. Dataset #Classes Image Size Train Val Test Total Train / class CIFAR-100 100 3×32×323× 32× 32 50,000 — 10,000 60,000 500 Tiny ImageNet 200 3×64×643× 64× 64 100,000 10,000 10,000 120,000 500 Notes. (1) CIFAR-10/100 provide no official validation split; a subset of the training set is commonly reserved as dev/val. (2) CIFAR-100 contains 100 fine-grained classes; 20 coarse superclasses are also defined for hierarchical labeling. (3) Tiny ImageNet is a subset of ImageNet synsets: per class 500 train, 50 val, and 50 test images (test labels are not publicly released). (4) All three datasets are single-label classification with RGB images resized to fixed resolutions. We evaluate our methods on two widely-used benchmark datasets in federated learning: CIFAR-100 and Tiny ImageNet. • CIFAR-100 (Krizhevsky et al., 2009): Contains 100 classes with 600 color images per class at a resolution of 32×3232× 32. It is a standard benchmark for evaluating federated image classification methods. • Tiny ImageNet: A subset of ImageNet with 200 classes and 500 images per class, providing a more challenging and high-resolution classification task. D.2 Federated Learning Configuration We simulate a cross-device federated learning environment using the following settings: Table 8: Hyperparameter configuration of ResNet-18 and Vit-Tiny (CIFAR100, Tiny-ImageNet) across different algorithms. Method Local Optimizer Local LR β β1 _1 β2 _2 Weight Decay FedAvg (Local SGD) SGD 0.1 — — — 0.001 SCAFFOLD SGD 0.1 — — — 0.001 FedCM SGD 0.1 0.9 — — 0.001 Local AdamW AdamW 3e-4 — 0.9 0.999 0.01 Local Sophia Sophia 3e-4 — 0.9 0.99 0.01 Local Muon Muon 3e-2 — 0.9 0.95 0.01 Local SOAP SOAP 3e-3 — 0.95 0.95 0.01 FedPAC_Sophia Muon 3e-4 0.5 0.9 0.99 0.01 FedPAC_Muon Muon 3e-2 0.5 0.9 0.95 0.01 FedPAC_SOAP SOAP 3e-3 0.5 0.95 0.95 0.01 Table 9: Hyperparameter configuration of ViT-Base fine-tuning across different algorithms. Method Local Optimizer Local LR β β1 _1 β2 _2 Weight Decay FedAvg (Local SGD) SGD 0.1 — — — 0.001 SCAFFOLD SGD 0.1 — — — 0.001 FedCM SGD 0.1 0.9 — — 0.001 Local AdamW AdamW 1e-4 — 0.9 0.999 0.01 Local Sophia Sophia 1e-4 — 0.9 0.99 0.01 Local Muon Muon 1e-2 — 0.9 0.95 0.01 Local SOAP SOAP 1e-3 — 0.95 0.95 0.01 FedPAC_Sophia Muon 1e-4 0.5 0.9 0.99 0.01 FedPAC_Muon Muon 1e-2 0.5 0.9 0.95 0.01 FedPAC_SOAP SOAP 1e-3 0.5 0.95 0.95 0.01 D.3 Model Architecture We adopt ResNet-18 as the backbone model. To better adapt it to CIFAR-100, we modify its architecture following standard practices (He et al., 2016): • Replace the original 7×77× 7 convolution with a 3×33× 3 kernel. • Remove the initial downsampling layers (stride-2 convolution and max pooling). We also compare Batch Normalization (BN) and Group Normalization (GN) in ResNet-18. Empirically, BN outperforms GN on CIFAR-100, so we adopt the BN-based version, denoted as ResNet-18-BN, throughout our experiments. D.4 Setting for ViT-Tiny We construct a lightweight Vision Transformer model, ViT-Tiny, specifically tailored for federated learning on the CIFAR-100 dataset. The design is based on the standard ViT architecture (Dosovitskiy et al., 2020), with modifications to accommodate the small input size and limited data per client. • Input resolution: 32×3232× 32 • Patch size: 4×44× 4, resulting in 64 tokens per image • Embedding dimension: 192 • Number of Transformer layers: 6 • Number of attention heads: 3 • Normalization: LayerNorm (applied before attention and MLP blocks) • Classification head: Linear projection to 100 classes (CIFAR-100) • Activation: GELU • Initialization: Xavier/Glorot for linear layers; sinusoidal positional encoding To regularize training, we apply dropout (0.1) to both attention and MLP layers. All models are trained from scratch without pretraining. Remarks. Due to the smaller capacity of ViT-Tiny and limited data per client, we find that careful normalization (e.g., LayerNorm placement) and early learning rate warmup are beneficial. For future work, more advanced token-mixing techniques or hybrid CNN-ViT backbones may further improve performance in federated settings. D.5 Transformer Fine-tuning Settings To demonstrate the effectiveness of our method on large-scale vision models, we conduct fine-tuning experiments using ViT-Base on Tiny ImageNet and CIFAR-100. For both models, we initialize from official ImageNet-22K pre-trained weights (Liu et al., 2021; Dosovitskiy et al., 2020) to ensure consistency across methods. We fine-tune all layers during federated training. Data Preprocessing. To align with the input resolution required by ViT, we resize images from both datasets to 224×224224× 224 using bilinear interpolation. Standard data augmentation techniques such as random cropping, horizontal flipping, and RandAugment are applied locally at the client side. D.6 Additional Federated Training Configuration of LLM To evaluate our algorithm under a smaller-scale federation, we further conduct experiments with a reduced number of clients and adjusted participation parameters. Federated Setup. We simulate a federated learning environment with the following configuration: Table 10: Hyperparameter configuration of LLAMA (C4) across different algorithms. Method Local Optimizer Local LR β β1 _1 β2 _2 Weight Decay FedAvg (Local SGD) SGD 0.1 — — — 0.001 SCAFFOLD SGD 0.1 — — — 0.001 FedCM SGD 0.1 0.9 — — 0.001 Local AdamW AdamW 3e-4 — 0.9 0.999 0.01 Local Sophia Sophia 3e-4 — 0.9 0.99 0.01 Local Muon Muon 3e-2 — 0.9 0.95 0.01 Local SOAP SOAP 3e-3 — 0.95 0.95 0.01 FedPAC_Sophia Muon 3e-4 0.5 0.9 0.99 0.01 FedPAC_Muon Muon 3e-2 0.5 0.9 0.95 0.01 FedPAC_SOAP SOAP 3e-3 0.5 0.95 0.95 0.01 Appendix E Appendix C: Experimental Appendix E.1 Communication and Computation Cost Analysis Table 11: Per-round communication cost of different momentum aggregation strategies. Here |x||x| denotes the number of model parameters (in floats), and CommCostCommCost is per-round communication cost, Compute-Cost is per-round computation time. (ViT-Tiny, R=300R=300, Dir-0.1, K=50K=50) Method / Strategy Communication CommCost Compute-Cost (s) Acc(%) FedAvg |x||x| 22.8 MB 4.56 s 27.24 SCAFFOLD 2|x|2|x| 45.6 MB 5.22 s 26.86 FedCM |x||x| 22.8 MB 4.68 s 16.95 Local AdamW |x||x| 22.8 MB 4.89 s 37.57 Local Sophia |x||x| 22.8 MB 4.92 s 34.05 Local Muon |x||x| 22.8 MB 5.14 s 44.00 Local SOAP |x||x| 22.8 MB 5.56 s 49.41 FedPAC_Sophia |x|+|Θ||x|+| | 45.6 MB 5.08 s 39.79 FedPAC_Muon |x|+1|Θ||x|+1| | 45.6 MB 5.25 s 47.81 FedPAC_SOAP |x|+2|Θ||x|+2| | 68.4 MB 5.68 s 51.16 FedPAC_Sophia_Light |x|+0.05|Θ||x|+0.05| | 23.9 MB 5.11 s 39.45 FedPAC_Muon_Light |x|+0.05|Θ||x|+0.05| | 23.9 MB 5.28 s 47.23 FedPAC_SOAP_Light |x|+0.1|Θ||x|+0.1| | 25.2 MB 5.69 s 50.56 Table 12: Test accuracy, training loss of each method on CIFAR-100 and Tiny-Imagenet using ResNet-18 over 300 communication rounds under IID Dir-0.5, Dir-0.1, Dir-0.05 (100 clients, 10% participation, batch size 50, K=50K=50). Method CIFAR-100 (ResNet-18) Tiny-Imagenet (ResNet-18) iid Dir-0.5 Dir-0.1 Dir-0.05 iid Dir-0.5 Dir-0.1 Dir-0.05 FedAvg 65.04 65.07 60.17 56.75 53.88 52.55 47.48 43.80 SCAFFOLD 66.04 65.51 60.69 56.43 54.49 53.36 47.76 43.92 FedCM 71.12 69.85 66.61 62.65 46.51 44.06 41.16 36.00 Local AdamW 64.37 63.58 59.23 55.24 51.06 49.97 44.01 40.00 Fed-Sophia 62.56 60.62 57.29 51.02 49.86 47.62 41.89 36.65 FedPM 62.86 61.23 57.25 50.98 49.63 48.21 41.53 36.62 Local Sophia 62.17 60.96 56.65 50.89 49.28 47.99 41.23 36.15 Local Muon 72.78 73.0273.02 67.26 49.86 60.5060.50 60.4060.40 52.83 34.76 Local SOAP 71.98 71.19 68.44 58.16 58.01 56.82 54.42 50.02 FedPAC_Sophia 64.90 64.71 59.96 53.66 49.30 50.92 43.81 36.37 FedPAC_Muon 72.7972.79 72.50 71.8571.85 65.5665.56 58.85 58.31 57.9557.95 54.0054.00 FedPAC_SOAP 69.79 71.05 69.25 64.16 56.57 56.22 55.62 51.81 Table 13: Test accuracy, training loss of each method on CIFAR-100 using ViT-Tiny over 300 communication rounds under IID, Dir-0.5, Dir-0.1, Dir-0.05 (100 clients, 10% participation, batch size 50, K=50K=50). Method CIFAR-100 (ViT-Tiny) Tiny-Imagenet (ViT-Tiny) iid Dir-0.5 Dir-0.1 Dir-0.05 iid Dir-0.5 Dir-0.1 Dir-0.05 FedAvg 32.97 33.66 27.24 23.42 18.61 18.03 15.68 14.05 SCAFFOLD 32.29 33.01 26.86 23.23 18.32 18.36 15.70 14.21 FedCM 21.97 22.02 16.95 14.74 10.23 11.66 8.88 8.15 Local AdamW 41.27 41.21 37.57 36.06 27.94 26.82 24.31 21.35 Fed-Sophia 38.34 38.56 34.23 32.27 23.56 24.46 23.85 22.17 FedPM 38.56.94 38.75 34.29 32.62 23.86 24.51 23.89 22.85 Local Sophia 37.94 38.48 34.05 32.25 22.90 23.46 22.49 21.14 Local Muon 50.50 51.75 47.81 41.76 34.80 34.57 30.51 28.25 Local SOAP 52.85 51.61 49.41 41.68 34.81 33.91 33.30 30.36 FedPAC_Sophia 42.53 42.29 39.79 32.71 23.26 22.16 23.01 22.37 FedPAC_Muon 52.70 52.02 44.00 39.68 31.78 33.48 31.45 30.25 FedPAC_SOAP 52.86 52.46 51.16 47.55 35.56 35.12 34.32 31.33 E.2 More baseline experiment comparisons Training on CIFAR-100 with ResNet-18. Table 12 reports the final test accuracy after 300 communication rounds on CIFAR-100 and Tiny-ImageNet with ResNet-18 under increasing data heterogeneity (IID, Dir-0.50.5, Dir-0.10.1, Dir-0.050.05; 100 clients, 10% participation, batch size 50, and K=50K=50 local steps). Overall, performance degrades as the Dirichlet concentration decreases, highlighting the non-trivial impact of client drift under severe non-IID distributions. Specialized second-order FL baselines. To directly address the concern that our gains may come from comparing against local second-order optimizers only, we additionally include two specialized second-order FL methods, Fed-Sophia and FedPM, which are designed to adapt curvature-aware (Sophia-style) updates to the federated setting. As shown in Table 12, both specialized baselines still exhibit noticeable performance drops as heterogeneity increases. For example, on CIFAR-100, Fed-Sophia decreases from 60.62 (Dir-0.50.5) to 51.02 (Dir-0.050.05), and FedPM decreases from 61.23 (Dir-0.50.5) to 50.98 (Dir-0.050.05), indicating that merely introducing a specialized second-order aggregation scheme does not fully resolve the instability induced by preconditioner drift. FedPAC improves robustness under strong heterogeneity. Across both datasets, FedPAC_Muon yields the best or near-best accuracy under strong non-IID settings, demonstrating substantial robustness benefits. On CIFAR-100 with Dir-0.050.05, FedPAC_Muon achieves 65.56, improving over FedAvg (56.75) by +8.81 and over Local Muon (49.86) by +15.70. On Tiny-ImageNet with Dir-0.050.05, FedPAC_Muon reaches 54.00, outperforming FedAvg (43.80) by +10.20 and Local Muon (34.76) by +19.24. These results align with our analysis that the server-side alignment and client-side correction in FedPAC explicitly mitigate the mismatch among client preconditioners, which becomes particularly harmful when data heterogeneity is severe. FedPAC also complements Sophia-style methods. Compared with the specialized Sophia-style FL baselines, FedPAC_Sophia provides consistent gains on CIFAR-100 across all heterogeneity levels (e.g., 53.66 vs. 51.02/50.98 under Dir-0.050.05 for Fed-Sophia/FedPM). On Tiny-ImageNet, FedPAC_Sophia is competitive with Fed-Sophia/FedPM and improves them under mild-to-moderate heterogeneity (IID/Dir-0.50.5/Dir-0.10.1), while being essentially on par under Dir-0.050.05. Taken together, these comparisons suggest that FedPAC addresses a more general failure mode—preconditioner drift—and thus can serve as a robust, optimizer-agnostic enhancement for federated second-order training. Takeaway. Table 12 verifies that (i) directly deploying local second-order optimizers can be brittle under strong heterogeneity, and (i) even specialized second-order FL methods may not fully eliminate the degradation, whereas FedPAC substantially improves robustness and accuracy, especially in the most heterogeneous regimes. Training on CIFAR-100 with ViT-Tiny. E.3 More baseline experiment on iid data To further strengthen the empirical comparisons (especially on IID data) and to examine the robustness trend as heterogeneity increases, we report additional results under IID and Dirichlet partitions with varying concentration parameters. Tables 12 and 13 summarize the final test accuracy after 300 communication rounds on CIFAR-100 and Tiny-ImageNet with ResNet-18 and ViT-Tiny, respectively (100 clients, 10% participation, batch size 50, K=50K=50). ResNet-18 results. From Table 12, we observe that under IID (and mild heterogeneity such as Dir-0.50.5), several strong baselines can be competitive (e.g., Local Muon on CIFAR-100). However, as heterogeneity becomes severe (Dir-0.10.1 and Dir-0.050.05), methods that rely on purely local second-order states (e.g., Local Muon) or specialized second-order FL baselines (e.g., Fed-Sophia and FedPM) degrade markedly, indicating the growing impact of preconditioner drift across clients. In contrast, FedPAC_Muon remains consistently strong and achieves the best accuracy under the most heterogeneous regime on both datasets, demonstrating improved robustness under non-IID distributions. ViT-Tiny results. Table 13 shows a similar trend for transformer-style backbones: while several baselines perform reasonably under IID/Dir-0.50.5, severe heterogeneity leads to substantial drops. Notably, FedPAC_SOAP achieves the best performance across all heterogeneity levels on both CIFAR-100 and Tiny-ImageNet with ViT-Tiny, suggesting that FedPAC can effectively stabilize curvature-aware training for modern architectures under client heterogeneity. Appendix F Related Work Family What is synchronized (server ↔ clients)? Extra cost (comm./state) Main issue it targets FedAvg / Local-SGD Model parameters/updates x only ≈|x|≈|x| per round General FL baseline; does not explicitly correct non-IID induced drift FedOpt (server adaptive; e.g., FedAdam/FedYogi) x only; optimizer state (moments) kept on server (not shared as a common client geometry) ≈|x|≈|x| (no extra client-side sync) Improves global update scaling/adaptivity, but does not align heterogeneous client metrics Control-variate / drift correction (e.g., SCAFFOLD) x + first-order control variates (gradient-level states) typically larger (e.g., SCAFFOLD uses ≈2|x|≈ 2|x| in our comm accounting) Corrects first-order client drift under non-IID via control variates; operates at gradient level, not metric level Naive second-order FL (Local Sophia/SOAP/Muon + FedAvg agg.) x only; each client maintains local preconditioner state Θ but does not synchronize it ≈|x|≈|x| per round Fails under non-IID due to preconditioner drift (mismatched curvature-induced geometries) FedPAC (ours) x + preconditioner state Θ (geometry / metric): server aggregates Θ and broadcasts a global reference (Alignment); local steps are corrected using a global preconditioned direction (Correction) ≈|x|+c|Θ|≈|x|+c| | (e.g., c=1c=1 for Muon/Sophia, c=2c=2 for SOAP; light variants use low-rank upload such as 0.05|Θ|0.05| |–0.1|Θ|0.1| |) Targets preconditioner drift directly by synchronizing the metric (Alignment) and suppressing long-term drift (Correction) Table 14: Positioning of FedPAC. Unlike FedOpt and control-variate methods that synchronize only parameter/first-order states, FedPAC explicitly synchronizes the curvature-defined geometry (preconditioner state Θ ) to mitigate preconditioner drift under non-IID data. Here |x||x| denotes per-round model communication in our accounting and |Θ|| | denotes the size of the optimizer preconditioner state. Proposition F.1 (Geometry drift induces preconditioned disagreement). Consider a second-order FL update where each client applies a preconditioner operator PΘP_ (defined by its optimizer state Θ ) to a local vector uiu_i (e.g., stochastic gradient). Assume PΘP_ is Lipschitz in the state: for any Θ,Θ′ , and any v, ‖PΘ(v)−PΘ′(v)‖≤LΘ‖Θ−Θ′‖‖v‖.\|P_ (v)-P_ (v)\|≤ L_ \,\| - \|\,\|v\|. (90) Define the per-round preconditioner drift metric ΔrD:=1S∑i∈Sr‖Θr,Ki−Θ¯r,K‖2. _r^D:= 1S _i∈ S_rE\| _r,K^i- _r,K\|^2. (91) Then for any collection uii∈Sr\u_i\_i∈ S_r, ‖1S∑i∈Sr(PΘr,ki(ui)−PΘ¯r,K(ui))‖2≤LΘ2(1S∑i∈Sr‖ui‖2)ΔrD.E \| 1S _i∈ S_r (P_ _r,k^i(u_i)-P_ _r,K(u_i) ) \|^2\;≤\;L_ ^2 ( 1S _i∈ S_rE\|u_i\|^2 ) _r^D. (92) In particular, if supi,r,k‖ui‖2≤G2 _i,r,kE\|u_i\|^2≤ G^2, then ‖1S∑i∈Sr(PΘr,ki(ui)−PΘ¯r,K(ui))‖2≤LΘ2G2ΔrD.E \| 1S _i∈ S_r (P_ _r,k^i(u_i)-P_ _r,K(u_i) ) \|^2\;≤\;L_ ^2G^2 _r^D. (93) Therefore, even when there is no gradient/client mismatch (e.g., ui≡u_i≡ u for all clients, or heterogeneity is negligible), heterogeneous optimizer states Θr,ki _r,k^i alone can induce a non-vanishing disagreement in the preconditioned update direction, which is a distinct error source from gradient-level client drift. Theorem F.2 (Geometry drift induces preconditioned disagreement). Consider a round r with a participating client set SrS_r of size S. Each client maintains a (second-order) optimizer state Θr,ki _r,k^i and applies the corresponding preconditioner operator PΘP_ to a local update vector uiu_i (e.g., stochastic gradient). Assume the operator is Lipschitz in the state: there exists LΘ>0L_ >0 such that for any Θ,Θ′ , and any v, ‖PΘ(v)−PΘ′(v)‖≤LΘ‖Θ−Θ′‖‖v‖.\|P_ (v)-P_ (v)\|≤ L_ \,\| - \|\,\|v\|. (94) Let Θ¯r,K:=1S∑i∈SrΘr,Ki _r,K:= 1S _i∈ S_r _r,K^i be the averaged state at the end of local steps, and define the preconditioner drift metric ΔrD:=1S∑i∈Sr‖Θr,Ki−Θ¯r,K‖2. _r^D:= 1S _i∈ S_rE\| _r,K^i- _r,K\|^2. (95) Then the disagreement between using heterogeneous local geometries and the averaged geometry satisfies ‖1S∑i∈Sr(PΘr,ki(ui)−PΘ¯r,K(ui))‖2≤LΘ2(1S∑i∈Sr‖ui‖2)ΔrD.E \| 1S _i∈ S_r (P_ _r,k^i(u_i)-P_ _r,K(u_i) ) \|^2\;≤\;L_ ^2 ( 1S _i∈ S_rE\|u_i\|^2 ) _r^D. (96) In particular, if supi,r,k‖ui‖2≤G2 _i,r,kE\|u_i\|^2≤ G^2, then ‖1S∑i∈Sr(PΘr,ki(ui)−PΘ¯r,K(ui))‖2≤LΘ2G2ΔrD.E \| 1S _i∈ S_r (P_ _r,k^i(u_i)-P_ _r,K(u_i) ) \|^2\;≤\;L_ ^2G^2\, _r^D. (97) Corollary F.3 (Drift penalty vanishes under geometry alignment). If the client geometries are aligned in round r in the sense that Θr,ki=Θ¯r,K _r,k^i= _r,K for all i∈Sri∈ S_r (equivalently, ΔrD=0 _r^D=0), then the preconditioned disagreement term in (96) is zero: 1S∑i∈SrPΘr,ki(ui)=1S∑i∈SrPΘ¯r,K(ui). 1S _i∈ S_rP_ _r,k^i(u_i)\;=\; 1S _i∈ S_rP_ _r,K(u_i). (98) Consequently, any residual inconsistency across clients is solely due to the vector-level mismatch in ui\u_i\ (e.g., gradient/client drift from non-IID data), rather than a mismatch in the preconditioning geometry. Remark F.4 (Why geometry drift is distinct from gradient/client drift). Gradient/client drift concerns the mismatch of the update vectors ui\u_i\ across clients (e.g., ∇fi≠∇f∇ f_i≠∇ f), and control-variate methods aim to reduce this vector-level discrepancy. In contrast, geometry drift concerns a mismatch of the operators PΘr,ki\P_ _r,k^i\, i.e., clients optimize under different local metrics defined by their second-order states. Theorem shows that even if ui\u_i\ were identical across clients (no gradient/client drift), heterogeneous geometries (ΔrD>0 _r^D>0) alone induce a non-vanishing disagreement in the preconditioned direction. Corollary further clarifies that synchronizing/aligning Θ removes this geometry-induced error source, leaving only the conventional gradient/client drift to be handled. Algorithm 4 Local SOAP Algorithm 0: Step of Local SOAP for an m×nm× n layer. Per layer, we maintain four matrices: L∈ℝm×mL ^m× m, R∈ℝn×nR ^n× n, V,M∈ℝm×nV,M ^m× n. Hyperparameters: learning rate η, betas (β1,β2)( _1, _2), epsilon ϵε, and preconditioning frequency f, communication rounds T, local updates K, the number of client N. 1: for t=0,…,Tt=0,…,T do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: for k=1,…,Kk=1,…,K do 4: Sample batch Bit,kB_i^t,k 5: Git,k∈ℝm×n←−∇WϕBt(Wit,k)G_i^t,k ^m× n←- _W _B_t(W_i^t,k) 6: End of gradient step, now update L and R and possibly also QLQ_L and QRQ_R 7: L←β2L+(1−β2)(GG⊤)L← _2L+(1- _2)(G ) 8: R←β2R+(1−β2)(G⊤G)R← _2R+(1- _2)(G G) 9: if tmodf=0t f=0 then 10: QL←Eigenvectors(L,QL)Q_L← Eigenvectors(L,Q_L) 11: QR←Eigenvectors(R,QR)Q_R← Eigenvectors(R,Q_R) 12: end if 13: git,k←QL⊤GQRg_i^t,k← Q_L GQ_R 14: 15: UPDATE(git,kg_i^t,k) by Adam 16: Mit,k←β1Mit,k+(1−β1)git,kM_i^t,k← _1M_i^t,k+(1- _1)g_i^t,k 17: Vit,k←β2Vit,k+(1−β2)(git,k⊙git,k)V_i^t,k← _2V_i^t,k+(1- _2)(g_i^t,k g_i^t,k) 18: Nit,k←Mit,k/(Vit,k+ϵ)N_i^t,k← M_i^t,k/( V_i^t,k+ε) 19: 20: Now that we have preconditioned by Adam in the rotated space, we go back to the original space 21: N~it,k←QLNit,kQR⊤ N_i^t,k← Q_LN_i^t,kQ_R 22: Wit,k←Wit,k−1−ηN~it,kW_i^t,k← W_i^t,k-1-η N_i^t,k 23: end for 24: Clint i communicate (Wit,K−Wit,0)(W_i^t,K-W_i^t,0) to Server; 25: end for 26: Wt+1←Wt+1N∑i(Wit,K−Wit,0)W^t+1← W^t+ 1N _i(W_i^t,K-W_i^t,0) 27: end for 28: 29: Eigenvectors(P,Q)(P,Q): 30: S←PQS← PQ 31: Q←QR(S)Q (S) Using QR decomposition 32: Return Q Algorithm 5 FedPAC_SOAP Algorithm 0: Step of FedPAC_SOAP for an m×nm× n layer. Per layer, we maintain four matrices: L∈ℝm×mL ^m× m, R∈ℝn×nR ^n× n, V,M∈ℝm×nV,M ^m× n. Hyperparameters: learning rate η, betas (β1,β2)( _1, _2), epsilon ϵε, and preconditioning frequency f, communication rounds R, local updates K, the number of client N. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: for k=1,…,Kk=1,…,K do 4: Sample batch Bir,kB_i^r,k 5: Gir,k∈ℝm×n←−∇Fi(xir,k;ξir,k)G_i^r,k ^m× n←-∇ F_i(x_i^r,k; _i^r,k) 6: End of gradient step, now update L and R and possibly also QLQ_L and QRQ_R 7: L←β2L+(1−β2)(GG⊤)L← _2L+(1- _2)(G ) 8: R←β2R+(1−β2)(G⊤G)R← _2R+(1- _2)(G G) 9: if tmodf=0t f=0 then 10: QL←Eigenvectors(L,QL)Q_L← Eigenvectors(L,Q_L) 11: QR←Eigenvectors(R,QR)Q_R← Eigenvectors(R,Q_R) 12: end if 13: git,k←QL⊤GQRg_i^t,k← Q_L GQ_R 14: 15: UPDATE(gir,kg_i^r,k) by Adam 16: Mir,k←β1Mir,k+(1−β1)gir,kM_i^r,k← _1M_i^r,k+(1- _1)g_i^r,k 17: Vir,k←β2Vir,k+(1−β2)(gir,k⊙gir,k)V_i^r,k← _2V_i^r,k+(1- _2)(g_i^r,k g_i^r,k) 18: Nir,k←Mir,k/(Vir,k+ϵ)N_i^r,k← M_i^r,k/( V_i^r,k+ε) 19: 20: Now that we have preconditioned by Adam in the rotated space, we go back to the original space 21: N~ir,k←QLNir,kQR⊤ N_i^r,k← Q_LN_i^r,kQ_R 22: ir,k+1=ir,k−η[(1−β)N~ir,k+βGr] x^r,k+1_i\!=\! x^r,k_i\!\!-\!η[(1\!-\!β) N_i^r,k\!+\!β _G^r]; 23: end for 24: Client i communicate (xir,K−xir,0)(x_i^r,K-x_i^r,0) to Server; 25: end for 26: Gr+1=−1SKη∑i=1S(ir,K−ir,0) _G^r+1=- 1SKη _i=1^S( x^r,K_i- x^r,0_i); 27: r+1=r+1S∑i=1S(ir,K−ir,0) x^r+1= x^r+ 1S _i=1^S( x^r,K_i- x^r,0_i); 28: end for 29: 30: Eigenvectors(P,Q)(P,Q): 31: S←PQS← PQ 32: Q←QR(S)Q (S) Using QR decomposition 33: Return Q Algorithm 6 Local Muon Algorithm (Federated Setting) 0: Local Muon for an m×nm× n weight matrix W (e.g., a Linear layer). Per layer, maintain one matrix: M∈ℝm×nM ^m× n (momentum). Hyperparameters: learning rate η, momentum β, weight decay λ (optional), epsilon ϵε, Newton–Schulz steps s (typically 55), dimension scaling γ(m,n)γ(m,n), communication rounds T, local steps K, number of clients N. 1: for t=0,…,Tt=0,…,T do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: Initialize local model: Wit,0←WtW_i^t,0← W^t 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bit,kB_i^t,k 6: Git,k∈ℝm×n←∇WϕBit,k(Wit,k−1)G_i^t,k ^m× n← _W _B_i^t,k(W_i^t,k-1) 7: Mit,k←β1Mit,k−1+(1−β1)Git,kM_i^t,k← _1\,M_i^t,k-1+(1- _1)\,G_i^t,k 8: Orthogonalize (typically on momentum) via Newton–Schulz 9: Uit,k←NewtonSchulz(Mit,k,s,ϵ)U_i^t,k← NewtonSchulz(M_i^t,k,s,ε) U is approximately the orthogonal factor of M 10: Apply dimension scaling (one common theoretical form: γ(m,n)=m/nγ(m,n)= m/n) 11: ΔWit,k←γ(m,n)Uit,k W_i^t,k←γ(m,n)\,U_i^t,k e.g. γ=fan-out/fan-inγ= fan-out/ fan-in 12: Update 13: Wit,k←Wit,k−1−η(ΔWit,k+λWit,k−1)W_i^t,k← W_i^t,k-1-η ( W_i^t,k+λ W_i^t,k-1 ) optional weight decay 14: end for 15: Client i sends (Wit,K−Wit,0)(W_i^t,K-W_i^t,0) to Server 16: end for 17: Wt+1←Wt+1N∑i=1N(Wit,K−Wit,0)W^t+1← W^t+ 1N _i=1^N(W_i^t,K-W_i^t,0) 18: end for 19: 20: NewtonSchulz(G,s,ϵ)(G,s,ε): 21: X←G/(‖G‖F+ϵ)X← G/(\|G\|_F+ε) 22: if m>nm>n then X←X⊤X← X end if 23: for j=1,…,sj=1,…,s do 24: A←XX⊤A← X 25: B←bA+cA2B← bA+cA^2 26: X←aX+BX← aX+BX 27: end for 28: if m>nm>n then X←X⊤X← X end if 29: return X common coefficients: a=3.4445,b=−4.7750,c=2.0315a=3.4445,\,b=-4.7750,\,c=2.0315 Algorithm 7 FedPAC_Muon Algorithm 0: Step of FedPAC_Muon for an m×nm× n layer. Per layer, maintain one matrix: M∈ℝm×nM ^m× n (momentum). Hyperparameters: learning rate η, momentum β, epsilon ϵε, Newton–Schulz steps s, (optional) weight decay λ, (optional) dimension scaling γ(m,n)γ(m,n), communication rounds R, local updates K, number of clients N, FedPAC mixing coefficient β∈[0,1]β∈[0,1]. Server maintains global direction Gr _G^r. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: for k=0,…,K−1k=0,…,K-1 do 4: Sample batch Bir,kB_i^r,k 5: Gir,k∈ℝm×n←−∇Fi(ir,k;ξir,k)G_i^r,k ^m× n←-∇ F_i( x_i^r,k; _i^r,k) 6: Mir,k+1←β1Mir,k+(1−β1)Gir,kM_i^r,k+1← _1\,M_i^r,k+(1- _1)\,G_i^r,k 7: Muon: orthogonalize the (momentum) update direction 8: Uir,k←NewtonSchulz(Mir,k+1,s,ϵ)U_i^r,k← NewtonSchulz(M_i^r,k+1,s,ε) 9: U~ir,k←γ(m,n)Uir,k U_i^r,k←γ(m,n)\,U_i^r,k optional scaling, e.g. γ=m/nγ= m/n 10: FedPAC local update with global direction mixing 11: ir,k+1=ir,k−η[(1−β)U~ir,k+βGr] x_i^r,k+1= x_i^r,k-η [(1-β) U_i^r,k+β _G^r ] optionally add +λir,k+λ x_i^r,k inside the bracket 12: end for 13: Client i communicates (ir,K−ir,0)( x_i^r,K- x_i^r,0) to Server; 14: end for 15: Gr+1=−1NKη∑i=1N(ir,K−ir,0) _G^r+1=- 1NKη _i=1^N ( x_i^r,K- x_i^r,0 ); 16: r+1=r+1N∑i=1N(ir,K−ir,0) x^r+1= x^r+ 1N _i=1^N ( x_i^r,K- x_i^r,0 ); 17: end for 18: 19: NewtonSchulz(G,s,ϵ)(G,s,ε): 20: X←G/(‖G‖F+ϵ)X← G/(\|G\|_F+ε) 21: if m>nm>n then X←X⊤X← X end if 22: for j=1,…,sj=1,…,s do 23: A←XX⊤A← X 24: X←12X(3I−A)X← 12X(3I-A) classic Newton–Schulz for (XX⊤)−1/2(X )^-1/2 25: end for 26: if m>nm>n then X←X⊤X← X end if 27: return X Algorithm 8 Local Sophia Algorithm (Federated Setting) 0: Local Sophia for an m×nm× n layer W. Per layer, maintain two matrices: M,H∈ℝm×nM,H ^m× n (momentum and Hessian-diagonal EMA). Hyperparameters: learning rate η, betas (β1,β2)( _1, _2), epsilon ϵε, clipping ρ, Hessian update frequency fhf_h, communication rounds T, local steps K, number of clients N, weight decay λ (optional). 1: for t=0,…,Tt=0,…,T do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: Initialize local model: Wit,0←WtW_i^t,0← W^t 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bit,kB_i^t,k 6: Git,k←∇WϕBit,k(Wit,k−1)G_i^t,k← _W _B_i^t,k(W_i^t,k-1) 7: Mit,k←β1Mit,k−1+(1−β1)Git,kM_i^t,k← _1M_i^t,k-1+(1- _1)G_i^t,k 8: Update H (diagonal Hessian estimate) occasionally 9: if ((tK+k)modfh)=0((tK+k) f_h)=0 then 10: Sample Rademacher noise Uit,k∈±1m×nU_i^t,k∈\± 1\^m× n 11: HVit,k←∇W2ϕBit,k(Wit,k−1)Uit,kHV_i^t,k← _W^2 _B_i^t,k(W_i^t,k-1)\,U_i^t,k HVP via auto-diff (Pearlmutter trick) 12: H^it,k←Uit,k⊙HVit,k H_i^t,k← U_i^t,k HV_i^t,k Unbiased diag(H) estimator 13: Hit,k←β2Hit,k−1+(1−β2)max⁡(H^it,k,0)H_i^t,k← _2H_i^t,k-1+(1- _2) ( H_i^t,k,0) 14: else 15: Hit,k←Hit,k−1H_i^t,k← H_i^t,k-1 16: end if 17: Dit,k←Mit,k⊘max⁡(Hit,k,ϵ)D_i^t,k← M_i^t,k (H_i^t,k,ε) element-wise divide 18: Dit,k←clip(Dit,k,−ρ,ρ)D_i^t,k← clip(D_i^t,k,-ρ,ρ) 19: Wit,k←Wit,k−1−η(Dit,k+λWit,k−1)W_i^t,k← W_i^t,k-1-η (D_i^t,k+λ W_i^t,k-1 ) optional weight decay 20: end for 21: Client i sends (Wit,K−Wit,0)(W_i^t,K-W_i^t,0) to Server 22: end for 23: Wt+1←Wt+1N∑i=1N(Wit,K−Wit,0)W^t+1← W^t+ 1N _i=1^N(W_i^t,K-W_i^t,0) 24: end for 25: 26: clip(X,a,b)(X,a,b): element-wise min⁡(max⁡(X,a),b) ( (X,a),b). Algorithm 9 PAC_Sophia Algorithm (Federated Setting) 0: PAC_Sophia for an m×nm× n layer W. Per layer, maintain two matrices: M,H∈ℝm×nM,H ^m× n (momentum and Hessian-diagonal EMA). Hyperparameters: learning rate η, betas (β1,β2)( _1, _2), epsilon ϵε, clipping ρ, Hessian update frequency fhf_h, communication rounds R, local steps K, number of clients N, PAC mixing coefficient β∈[0,1]β∈[0,1], weight decay λ (optional). Server maintains global direction Gr _G^r. 1: for r=0,…,Rr=0,…,R do 2: for each client i∈1,…,Ni∈\1,…,N\ in parallel do 3: Initialize local model: Wir,0←WrW_i^r,0← W^r 4: for k=1,…,Kk=1,…,K do 5: Sample batch Bir,kB_i^r,k 6: Gir,k←∇WϕBir,k(Wir,k−1)G_i^r,k← _W _B_i^r,k(W_i^r,k-1) 7: Mir,k←β1Mir,k−1+(1−β1)Gir,kM_i^r,k← _1M_i^r,k-1+(1- _1)G_i^r,k 8: Update H (diagonal Hessian estimate) occasionally 9: if ((rK+k)modfh)=0((rK+k) f_h)=0 then 10: Sample Rademacher noise Uir,k∈±1m×nU_i^r,k∈\± 1\^m× n 11: HVir,k←∇W2ϕBir,k(Wir,k−1)Uir,kHV_i^r,k← _W^2 _B_i^r,k(W_i^r,k-1)\,U_i^r,k HVP via auto-diff (Pearlmutter trick) 12: H^ir,k←Uir,k⊙HVir,k H_i^r,k← U_i^r,k HV_i^r,k Unbiased diag(H) estimator 13: Hir,k←β2Hir,k−1+(1−β2)max⁡(H^ir,k,0)H_i^r,k← _2H_i^r,k-1+(1- _2) ( H_i^r,k,0) 14: else 15: Hir,k←Hir,k−1H_i^r,k← H_i^r,k-1 16: end if 17: Dir,k←Mir,k⊘max⁡(Hir,k,ϵ)D_i^r,k← M_i^r,k (H_i^r,k,ε) 18: Dir,k←clip(Dir,k,−ρ,ρ)D_i^r,k← clip(D_i^r,k,-ρ,ρ) Sophia direction 19: PAC mixing with global direction 20: Wir,k←Wir,k−1−η[(1−β)Dir,k+βGr+λWir,k−1]W_i^r,k← W_i^r,k-1-η [(1-β)D_i^r,k+β _G^r+λ W_i^r,k-1 ] λ optional 21: end for 22: Client i communicates (Wir,K−Wir,0)(W_i^r,K-W_i^r,0) to Server 23: end for 24: Gr+1=−1NKη∑i=1N(Wir,K−Wir,0) _G^r+1=- 1NKη _i=1^N (W_i^r,K-W_i^r,0 ); 25: Wr+1=Wr+1N∑i=1N(Wir,K−Wir,0)W^r+1=W^r+ 1N _i=1^N (W_i^r,K-W_i^r,0 ); 26: end for 27: 28: clip(X,a,b)(X,a,b): element-wise min⁡(max⁡(X,a),b) ( (X,a),b).

Instant research discovery