Conceptio › Archive › arXiv CS

arXiv CSopen access

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

arXiv:2604.25467v1 [cs.LG] 28 Apr 2026

Shuchen Zhu Peking University

Zhengyang Huang Beihang University

Yuqi Xu Peking University

Peijin Li Peking University

April 29, 2026 Abstract Federated learning increasingly operates in a large-model regime where communication, memory, and computation are all scarce. Typically, non-IID client data induce drift that degrades the stability and performance of local training. Existing remedies such as SCAFFOLD introduce heterogeneitycorrection mechanisms to address this challenge, but they incur substantial extra communication and memory overhead. This paper proposes a subspace optimization method for federated learning (SSF), which performs heterogeneity-corrected optimization in a low-dimensional subspace using only projected quantities, while preserving full-dimensional control information through a backfill-style update that retains residual components whenever the active subspace changes. Under standard smoothness and bounded√ e variance assumptions, SSF attains a non-asymptotic rate of order O(1/T + 1/ N KT ). Experiments show favorable accuracy–efficiency trade-offs under heterogeneous data.

Introduction

Federated learning (FL) has emerged as a central paradigm for privacy-preserving distributed training, where a central server coordinates optimization of a shared model across many clients that keep their data local. The standard Federated Averaging (FedAvg) algorithm amortizes communication by allowing multiple local updates between synchronization rounds, and has become a widely adopted baseline due to its simplicity and empirical success in practical deployments [16, 11, 21, 17]. However, modern FL applications increasingly involve models with hundreds of millions or even billions of parameters and large-scale datasets deployed on resource-constrained edge devices [32, 18]. In this large-model regime, federated optimization faces a triple challenge in terms of computation, memory, and communication: each local training step can be computationally heavy, storing full model states stresses limited device memory, and transmitting full-dimensional updates over constrained bandwidth becomes prohibitive. These challenges are further exacerbated in realistic heterogeneous environments where client data distributions differ substantially from each other and from the global distribution. In such non-IID settings, the local objective functions Fi become misaligned, inducing client drift in local updates and degrading the stability and convergence of FedAvg and its variants [12, 27, 26]. Heterogeneity-correction methods such as SCAFFOLD address this issue by introducing control variates on the server and clients to correct local gradients, leading to convergence guarantees that resemble centralized training even under strong heterogeneity and partial participation [9]. While highly effective as robustness baselines in small- to medium-scale settings, these methods are especially costly in the large-model regime: they require storing full-dimensional auxiliary states at each client, communicating additional full-dimensional control variates between the server and clients, and computing full-dimensional gradients at every local step. As a result, full-dimensional heterogeneitycorrection methods are at odds with the triple-efficiency requirements of large-scale FL on memory- and bandwidth-limited devices. To alleviate these resource constraints, several complementary lines of work have been explored in federated and distributed optimization. A first line focuses on reducing communication cost via gradient or model compression. Quantization and sparsification techniques such as QSGD and sparse SGD compress updates 1

before transmission, and error-feedback mechanisms compensate for the bias introduced by compression [1, 13, 25, 22, 23, 24, 5, 6, 7, 3]. These methods can substantially lower the number of bits communicated per round and are widely used in practice. However, they typically operate on full-dimensional gradients or model differences, so the communication cost still scales with the ambient dimension d of the model parameters, and the reduction is mainly in constant factors rather than dimensionality. More critically for large models, error-feedback schemes maintain full-dimensional error buffers to accumulate discarded information, which do not reduce—and often increase—the memory footprint on each client. When combined with full-dimensional control variates or other auxiliary states, compression and error-feedback can therefore aggravate memory pressure even as they reduce raw bandwidth usage. A second line of work leverages subspace and low-rank parameterizations to obtain more holistic efficiency gains. In centralized training, methods such as GaLore project gradients onto low-rank subspaces to reduce the cost of backpropagation while preserving performance [33]. In federated learning, subspace and lowrank approaches—such as those using low-dimensional adapters, partial model training, or low-rank drift variables in subspace-based federated learning methods [31, 14]—constrain optimization to a low-dimensional subspace or factorization. This design yields simultaneous reductions in computation, communication, and memory: local updates are computed in an r-dimensional space with r ≪ d, communicated quantities are low-dimensional, and only subspace parameters or low-rank factors need to be stored. However, in strongly heterogeneous FL, these methods often lack robust mechanisms to control client drift. Many subspace FL algorithms do not integrate heterogeneity correction, and naive attempts to introduce low-rank drift compensation variables face serious stability issues: when the underlying subspace frequently changes or rotates, the low-rank compensation variables must be repeatedly re-aligned, making it difficult to maintain effective updates and leading to degraded correction quality and convergence. Consequently, while subspace and low-rank methods excel in efficiency, their convergence robustness under heterogeneity remains limited. More broadly, existing approaches still exhibit a fundamental tension between robustness and efficiency along the three core resource dimensions. Full-dimensional correction and gradient-tracking methods maintain strong control of client drift but incur significant computation, communication, and memory overhead by operating in Rd [9, 10, 15]. Compression with error feedback can mitigate communication costs, yet continues to rely on full-dimensional auxiliary buffers, leaving the memory footprint large or even larger [22, 1, 13, 25, 3, 23, 7]. Subspace and low-rank FL methods reduce computation, communication, and memory by working in low-dimensional parameterizations [33, 31, 14], but they do not implement heterogeneitycorrection in a principled way. Naive combinations that keep static full-dimensional control variates while training in evolving subspaces either retain the full-dimensional auxiliary states—negating the subspace efficiency benefits—or discard accumulated correction information whenever the subspace changes, thereby undermining the effect of heterogeneity-correction. This state of affairs highlights a key gap: there is currently no federated optimization algorithm that performs native heterogeneity correction in a low-dimensional subspace, so as to (i) preserve robustness to data heterogeneity, while (ii) simultaneously achieving subspace-level efficiency in computation, communication, and memory. Our goal in this work is to fill this gap. Our contributions are as follows. • We propose SSF, a federated optimization algorithm that combines heterogeneity correction with subspace training to address the triple challenge of computation, memory, and communication in large-model FL. SSF constrains all optimization steps to a low-dimensional subspace defined by a projection matrix or low-rank parameterization: local updates and all communicated quantities live in this r-dimensional space, with r ≪ d, while the underlying goal remains to optimize the original highdimensional model parameters for downstream tasks. Crucially, SSF maintains server and client control variates natively in the optimization subspace and uses their difference to correct projected gradients during local updates, mirroring the SCAFFOLD correction mechanism without ever introducing fulldimensional auxiliary states. When the subspace basis is updated or rotated, SSF employs a translation mechanism that consistently maps the accumulated control variates into the new basis, preserving historical heterogeneity-correction information rather than discarding it when the subspace evolves. Through this subspace-native design, SSF achieves simultaneous savings in computation, communication, and memory, while explicitly correcting client drift under data heterogeneity. At the same time, SSF preserves the convergence behavior one would expect from single-node subspace training, so that the

efficiency gains from low-dimensional optimization do not come at the cost of degraded convergence rates. • We provide a comprehensive non-asymptotic convergence analysis for SSF under standard smoothness and bounded-variance assumptions on stochastic gradients. We first show that when gradients are projected through any matrix with spectral norm at most one and an appropriate local learning rate is used, the projected stochastic gradients satisfy the same type of bounded-variance property as in the full space. Building on a subspace client-drift bound and a harmonic coupling of the global and √ e local stepsizes, we obtain a transient iteration complexity of order O(1/T + 1/ N KT ) for the average squared gradient norm, where N is the number of participating clients per round and K is the number of local steps. This establishes linear speedup in N and K and shows that subspace-based heterogeneity correction can achieve SCAFFOLD-level convergence behavior while avoiding full-dimensional auxiliary states. • We empirically evaluate SSF in representative heterogeneous federated learning scenarios with large models and resource-constrained clients, comparing against full-dimensional heterogeneity-correction methods, compressed communication baselines, and existing subspace FL algorithms. The experiments demonstrate that SSF simultaneously reduces memory usage, communication volume, and local computation cost, while preserving robustness to data heterogeneity and achieving fast convergence. These empirical results confirm that the proposed subspace-based heterogeneity-correction mechanism can deliver strong accuracy–efficiency trade-offs in practical large-model FL deployments.

Related Work

The existing literature can be organized around three complementary but only partially compatible design goals: robustness to heterogeneity through variance reduction, communication efficiency through compression, and holistic efficiency through subspace or low-rank optimization. This taxonomy, which is strongly supported by the survey assets, is useful for positioning our work because Subspace-SCAFFOLD is intended to operate precisely at the intersection of these three lines rather than within any single one of them.

2.1

Federated Learning and Data Heterogeneity

The foundational Federated Averaging (FedAvg) algorithm established the paradigm of communicationefficient training via iterative local SGD and periodic model averaging [16, 11]. Its empirical success spurred extensive theoretical analysis, establishing convergence guarantees in convex and nonconvex settings and clarifying the communication–optimization trade-off induced by local steps [21, 12, 17, 27]. At the same time, these analyses make clear that statistical heterogeneity induces client drift: local iterates move toward client-specific objectives and can become systematically misaligned with the global objective, especially under partial participation or unreliable communication [12, 26]. Specialized remedies based on architecture or overparameterization can mitigate this phenomenon in certain settings [28, 8], but they do not provide a general optimization-level correction mechanism. A more principled response is heterogeneity-correction through auxiliary states. The SCAFFOLD algorithm introduces client and server control variates to correct the local descent direction, thereby reducing drift and recovering convergence guarantees that resemble centralized training even under strong heterogeneity and client sampling [9]. This perspective is closely related to gradient tracking in decentralized optimization [10, 15]. Subsequent accelerated, composite, and sharpness-aware variants further reinforce the value of correction terms in heterogeneous federated optimization [29, 2, 34, 30, 4]. The central limitation, however, is shared across this line of work: the correction states are full-dimensional. Each client must store and often communicate vectors in Rd , which creates a substantial memory and bandwidth burden for modern large models [9, 32, 18]. Our paper is motivated by preserving the robustness of this line while removing its full-dimensional systems cost.

2.2

Communication Compression and Error Feedback

A direct route to communication efficiency is to compress model updates using quantization or sparsification [1, 13, 25]. These primitives can greatly reduce bits per round, but aggressive compressors also introduce variance or deterministic bias. The now-standard remedy is error feedback, which stores the discrepancy between the intended update and its compressed version in a local residual buffer so that the lost information is gradually reinjected into future rounds [22]. This principle underlies compressed local-SGD frameworks and doubly compressed schemes such as Qsparse-local-SGD and DoubleSqueeze [3, 23]. It has also been strengthened by analyses showing that momentum can further improve error-feedback behavior [6]. In federated learning, compression has been integrated with heterogeneity-aware optimization in methods such as SCAFCOM and SCALLION, which combine controlled averaging with compressed communication and obtain convergence guarantees under arbitrary heterogeneity [7]. Adaptive compression strategies further tune this trade-off online [24]. Nevertheless, this literature inherits a key structural limitation: error feedback requires a full-dimensional residual buffer on each client, so the memory footprint still scales with the ambient model dimension [22, 3, 7]. In that sense, compression often shifts the bottleneck from raw communication to device memory rather than resolving both simultaneously. Moreover, under non-IID data, local residuals can themselves evolve heterogeneously, complicating their interaction with client drift. For our purposes, the main lesson from this line is that reducing transmitted bits alone is not enough; one must also control the dimension and persistence of the auxiliary state.

2.3

Subspace and Low-Rank Optimization Methods

A more holistic strategy is to reduce the dimension of the optimization process itself. In centralized largemodel training, low-rank projection methods such as GaLore show that gradients and optimizer states can be maintained in low-dimensional subspaces while preserving strong empirical performance [33]. Related survey materials emphasize a broader geometric viewpoint in which projection, preconditioning, and low-rank structure jointly reshape the optimization problem so as to trade off computation, communication, and memory more effectively. In federated settings, this idea gives rise to subspace and low-rank algorithms that project client updates, communicate only low-dimensional coordinates, or optimize within low-rank parameterizations [31, 14, 20]. Greedy low-rank compression methods in distributed optimization further illustrate that structured low-rank communication can admit convergence guarantees when combined with appropriate correction mechanisms [5]. The advantage of this line is triple efficiency: local computation is reduced because updates are formed in an r-dimensional space, communication is reduced because only projected quantities are transmitted, and memory is reduced because optimizer or auxiliary states can be maintained in low-dimensional coordinates. However, the survey assets also make clear that existing subspace methods generally prioritize efficiency over heterogeneity robustness. Some methods employ dual variables or heuristic correctors, but they do not furnish a SCAFFOLD-style, theoretically grounded heterogeneity-correction mechanism in the same evolving low-dimensional space [31]. This is the key incompatibility. Classical control variates are static and full-dimensional, whereas practical subspace methods often rely on dynamic or periodically refreshed bases. When the active subspace changes, a naive full-dimensional correction is either too costly to preserve or becomes misaligned with the new coordinates. Thus, the historical information required for effective drift correction is either retained at full cost or lost during subspace transitions.

2.4

Positioning of This Work

Our work is positioned exactly in the gap exposed by the above taxonomy. Heterogeneity-correction methods offer robustness but pay full-dimensional memory and communication costs; compression methods reduce transmitted bits but still rely on full-dimensional auxiliary buffers; and subspace methods offer holistic efficiency but lack a native, theoretically principled mechanism for correcting client drift under heterogeneity. The objective of Subspace-SCAFFOLD is to bridge these lines by co-designing subspace optimization and heterogeneity correction so that heterogeneity-correction information is maintained consistently with the lowdimensional optimization geometry rather than appended as an incompatible full-dimensional afterthought. In this sense, our work is not merely another compressed or low-rank FL method: it aims to recover the

robustness guarantees associated with SCAFFOLD while preserving the subspace efficiency required in largemodel federated learning. The theoretical development is given in Section 4, and the empirical consequences of this design are evaluated in Section 5.

Methodology

This section presents the methodology of SSF, including the problem setup, the algorithmic intuition, the main procedure, and the key implementation and complexity considerations. Our goal is to reduce the communication dimension from d to r while retaining the heterogeneity-correction mechanism through projected updates and a residual-preserving backfill step in the ambient space.

3.1

Problem setup and motivation

We consider the following stochastic federated optimization problem over N clients. N

F (x) =

1 X Fi (x), N i=1

(1)

where the variable x ∈ Rd , Fi = Eξ∼Di [fi (x, ξ)] and Di is the local dataset at client i. The dataset Di are heterogeneous across clients, leading to client drift when performing local gradient descent. The methodological motivation is straightforward. Standard FedAvg aggregates local model changes but does not explicitly correct the discrepancy between local and global gradient information. SCAFFOLD addresses this issue by maintaining global and local control variates, but its memory and communication cost scales with the ambient dimension. SSF seeks to preserve the same correction mechanism while restricting communication and active-device state transfer to a lower-dimensional random subspace.

3.2

Subspace projection and backfill mechanism

Let Pt ∈ Rr×d be a shared orthonormal projection operator at round t, with Pt Pt⊤ = Ir ,

(2)

and let Pt := Pt⊤ Pt d

(3) d

be the corresponding orthogonal projector in R . For any full-space vector x ∈ R , we decompose it into its projected and residual parts: xproj = Pt x,

xres = x − Pt⊤ xproj = (I − Pt )x.

(4)

The same decomposition is applied to the global control variate ct and to each client control variate cti . The key modeling idea is that optimization work is done only on the projected coordinates. Local iterates evolve in Rr , gradients are projected through Pt , and server aggregation also takes place in the subspace. After the projected update is computed, the method restores a full-space model by backfilling the preserved residual part: t xt+1 = Pt⊤ xt+1 (5) proj + xres . This backfill mechanism is one of the defining ingredients of SSF: it enables subspace communication without discarding the ambient-space representation of the iterate.

3.3

Algorithmic intuition

The algorithmic design follows two complementary principles. First, as in SCAFFOLD, each client performs locally corrected stochastic-gradient steps. In full rank, the correction takes the form gi − cti + ct , where cti tracks the local gradient and ct tracks the global gradient. In SSF, the same correction is applied after projection into the active subspace. This suppresses client drift in 5

the communicated directions while keeping the local and global correction terms dimension-compatible with the subspace iterate. Second, unlike purely subspace-dual methods, SSF preserves a full-dimensional control state. Only the current subspace component is refreshed using new gradient information, whereas the orthogonal complement is carried forward unchanged. Concretely, the global control update has the form N −1 K−1

ct+1 =

1 X X ⊤ Pt Pt git,k + ctres , N K i=0

(6)

k=0

with the analogous client update K−1

ct+1 = (I − Pt⊤ Pt )cti + i

1 ⊤ X t,k P Pt gi . K t

(7)

k=0

Hence, only the projected component is replaced, whereas the residual component is preserved. This residualpreserving update allows SSF to accumulate dual information across multiple rounds and multiple subspaces, rather than confining itself to a fixed low-rank dual representation. This methodology rests on three ingredients: a shared random subspace projection, projected control variables updates, and a backfill mechanism that restores a full-space iterate after each subspace update. In contrast, FedSub keeps only subspace dual variables and may therefore lose information when the projector changes, whereas SSF stores the control state in the full space and communicates only its projected component. This distinction is central to the implementation rationale.

3.4

Main algorithm

For completeness, we present the core algorithm analyzed in this paper. Algorithm 1 SSF 0 Require: Initial model x0 , control variables {c0i }N i=1 , c ; learning rates ηl , ηg ; local steps K; clients S; subspace dim. r; rounds T Initialize subspace projector: obtain P0 ∈ O(Rr×d ) by common random seed for t = 0, · · · , T − 1 do Decompose model and control: xtproj = Pt xt , xtres = xt −Pt⊤ xtproj , ctproj = Pt ct , ctres = ct −Pt⊤ ctproj for all client i in parallel do Client control split: cti,proj = Pt cti , cti,res = cti − Pt⊤ cti,proj t,0 Initialize yi,proj = xtproj for k = 0, . . . , K − 1 do t,k Reconstruct full model: yit,k = Pt⊤ yi,proj + xtres Compute stochastic gradient: git,k = ∇fi (yit,k ; ξit,k ) t,k+1 t,k Update projected model: yi,proj = yi,proj − ηl (Pt git,k − cti,proj + ctproj ) end for PK−1 1 Update client control: ct+1 = (I − Pt⊤ Pt )cti + K Pt⊤ Pt k=0 git,k i end for t,K t ηl K P xproj −yi,proj t Update projected model: xt+1 proj = xproj − ηg N i ηl K PN −1 PK−1 ⊤ t,k t t t+1 Backfill to full space: xt+1 = Pt⊤ xt+1 = N1K i=0 proj + xres , c k=0 Pt Pt gi + cres Update projector: generate new Pt+1 for next round end for

3.5

Pseudocode-level implementation view

At a schematic level, each round consists of the following four stages.

Stage 1: projector generation and decomposition. A common random seed generates the shared projector Pt . The server decomposes the full model and control variables into projected and residual components. Only the projected pieces are prepared for communication. Stage 2: local projected updates. For each sampled client, the server first splits the local control variable into projected and residual parts. The client initializes a projected iterate from the broadcast subspace model, reconstructs the corresponding full model at each local step, computes a stochastic gradient in the ambient space, projects that gradient back into the subspace, and performs a heterogeneity-corrected projected update. Stage 3: projected aggregation. After K local steps, the server aggregates the projected client endpoints exactly in the FedAvg style but entirely inside the r-dimensional subspace. Stage 4: full-space backfill and control refresh. The server lifts the projected model back to the ambient space by combining the new projected coordinates with the old residual. The global and client control variates are then refreshed only on the active subspace and keep their orthogonal complements unchanged. The implementation notes also stress a practical distinction between SSF and FedSub. FedSub updates and rotates only subspace dual variables, while SSF stores full-dimensional controls on the CPU and computes low-dimensional projections for communication. This design is heavier than FedSub in raw storage, but it avoids the structural information loss caused by repeated subspace-only dual rotations.

3.6

Implementation considerations

Several implementation choices are important for a robust realization of SSF. Shared projector generation. The projector is constructed from a common seed, often by sampling a Gaussian random matrix and orthonormalizing it via QR decomposition. This guarantees that all clients use the same subspace in a given round without explicitly transmitting a dense random basis every time. Projected communication only. The communicated objects are the subspace coordinates, such as xtproj , projected local endpoints, and projected control terms. This reduces the communication payload from order d to order r per vector-valued quantity. Full-space control storage with low-dimensional transfer. The full control variables can be stored on the CPU. For each round, the server computes their projected versions and transfers only those projected quantities to the GPU. The orthogonal residual components remain on CPU memory and are accessed only when the projector changes or when the full-space control variable is explicitly reconstructed. Projector refresh. The methodology allows the projector to change over time. In the baseline algorithm analyzed here this happens every round, although the refresh frequency can also be treated as a tunable implementation parameter. Refreshing more often increases directional coverage in the ambient space, while refreshing less often may reduce overhead and stabilize short-term local progress within a fixed subspace.

3.7

Complexity and practical trade-offs

Computation. SSF reduces computation, as the computation overhead of local stochastic-gradient is roughly proportional to the subspace dimension in most problems. Memory. Compared to Full-dimensional heterogeneity correction methods such as SCAFFOLD that store a full-dimensional control variables, SSF only requires its subspace low-dimensional version at local steps, significantly reducing the memory footprint of the client control variables. SSF still maintains full-dimensional control variables, but they can remain on CPU memory; the GPU needs only the projected low-dimensional versions. FedSub is lighter because its dual and endpoint states are only Θ(N rm), but that reduction comes with the information-loss mechanism discussed above. 7

Communication. The main systems advantage of SSF is that per-round communication scales with the subspace dimension. The effective communication cost is reduced by a factor on the order of r/d. Trade-off between compression and fidelity. Smaller r yields larger communication savings but can weaken the evolution of full-space stationarity, while larger r improves optimization fidelity at higher communication cost. SSF is designed to navigate this trade-off more robustly than purely subspace-dual alternatives because it preserves a full-space control memory across rounds.

Theory

This section presents the theoretical analysis of SSF. We begin with the standing assumptions and notation, and then give the full sequence of lemmas, the main convergence theorem, and the harmonic-stepsize corollary. The theorem, lemma, corollary, and proof materials are included in full, with labels preserved exactly.

4.1

Assumptions

Assumption 1 (L-smoothness). Each client loss Fi : Rd → R satisfies ∥∇Fi (x) − ∇Fi (y)∥ ≤ L∥x − y∥ for all x, y. Assumption 2 (Bounded variance and independence across samples). For stochastic gradients gi (x; ξ), E[gi (x; ξ) | x] = ∇Fi (x) and E∥gi (x; ξ) − ∇Fi (x)∥2 ≤ σ 2 . In addition, all stochastic gradient noises gi (x; ξ) − ∇Fi (x) are independent across different clients i. Assumption 3 (Shared random subspace). In each outer round k, all clients use the same orthonormal Pk ∈ O(Rr×d ), with Pk Pk⊤ = Ir and EPk [Pk⊤ Pk ] = (r/d)Id .

4.2

Notation and projector geometry

We summarize the notation that will be used throughout the analysis. Dimensions and subspace ratio. We denote by d the ambient dimension and by r the dimension of the random working subspace. The corresponding subspace ratio is ρ := dr ∈ (0, 1]. Random subspace projector. In each outer round t, all clients share the same random orthonormal matrix Pt ∈ Rr×d as in Assumption 3, with Pt Pt⊤ = Ir and EPt [Pt⊤ Pt ] = (r/d)Id . We write Pt := Pt⊤ Pt ∈ Rd×d for the associated orthogonal projector onto the random r-dimensional subspace. Unless otherwise stated, expectations E[·] are taken with respect to the joint randomness of the projector Pt , the sampled clients, and all stochastic gradients in the corresponding round.

4.3

Gradient and variance bounds

Lemma 1 (Gradient and Variance Bounds). Suppose Assumption 2 holds. Let gi (x; ξ) be the stochastic gradient computed by client i at model parameter x with random seed ξ. Then the following properties hold: 1. Unbiasedness: E[gi (x; ξ) | x] = ∇Fi (x). 2. Bounded Variance: E[∥gi (x; ξ) − ∇Fi (x)∥2 | x] ≤ σ 2 . 3. Bounded Second Moment: E[∥gi (x; ξ)∥2 | x] ≤ ∥∇Fi (x)∥2 + σ 2 . 4. Projected Variance: For any fixed matrix P (or random matrix P independent of ξ) with spectral norm ∥P ∥2 ≤ 1, E[∥P (gi (x; ξ) − ∇Fi (x))∥2 | x, P ] ≤ σ 2 .

Proof. The first two properties are restatements of Assumption 2. For the third property (Bounded Second Moment), we apply the bias-variance decomposition: E[∥gi (x; ξ)∥2 | x] = E[∥gi (x; ξ) − ∇Fi (x) + ∇Fi (x)∥2 | x] = E[∥gi (x; ξ) − ∇Fi (x)∥2 | x] + ∥∇Fi (x)∥2 + 2⟨E[gi (x; ξ) − ∇Fi (x) | x], ∇Fi (x)⟩ = E[∥gi (x; ξ) − ∇Fi (x)∥2 | x] + ∥∇Fi (x)∥2 2

(since the cross term is zero)

≤ σ + ∥∇Fi (x)∥ . For the fourth property (Projected Variance), let v = gi (x; ξ) − ∇Fi (x). Conditioned on x and P , we have: E[∥P v∥2 | x, P ] ≤ E[∥P ∥22 ∥v∥2 | x, P ] ≤ 12 · E[∥v∥2 | x] ≤ σ 2 . This concludes the proof.

4.4

Client drift bound in subspace

Lemma 2 (Client Drift Bound in Subspace). Suppose Assumptions 1 (L-smoothness) and 2 (bounded variance 1 and independence across samples and local steps) hold, and the local learning rate satisfies ηl ≤ 2KL . For any client i and outer round t, define K−1 X 2 t,k E yi,proj − xtproj . Ei,t := k=0

Then Ei,t ≤ 6K 2 ηl2 σ 2 + 12K 3 ηl2 ∇Fi (xt ) − cti + ct

(8)

In particular, the noise-induced part of the client drift scales as O(K 2 ηl2 σ 2 ) thanks to the independence of stochastic gradient noises across local steps for the same client within a round (see Assumption 2). Proof. Fix a round t and a client i, and write P = Pt for brevity. Recall the projected local iterates t,k t,0 yi,proj ∈ Rr with initialization yi,proj = xtproj , and define t,k t ∆t,k i := yi,proj − xproj ,

k = 0, 1, . . . , K.

Then ∆t,0 i = 0, and the projected update is t,k+1 t,k yi,proj = yi,proj − ηl P git,k − cti,proj + ctproj , t,k where git,k = ∇Fi (yit,k ; ξit,k ) and yit,k = P ⊤ yi,proj + xtres .

Step 1: Decomposition into deterministic drift and noise. Introduce the single-step gradient noise ξit,k := git,k − ∇Fi (yit,k ), so that git,k = ∇Fi (yit,k ) + ξit,k and, by Assumption 2, E[ξit,k | yit,k ] = 0,

E[∥ξit,k ∥2 | yit,k ] ≤ σ 2 .

Define the corrected direction t,k t t dt,k i := ∇Fi (yi ) − ci + c .

The update for ∆t,k becomes i

t,k t,k ∆t,k+1 = ∆t,k i i − ηl P di − ηl P ξi .

Unrolling from k = 0 and using ∆t,0 i = 0 yields ∆t,k+1 = −ηl i

k X

P dt,j i − ηl

j=0

k X j=0

P ξit,j .

(9)

Step 2: Basic quadratic split.

From (9) and ∥a + b∥2 ≤ 2∥a∥2 + 2∥b∥2 , we have

E ∆t,k+1 i

≤ 2ηl2 E

k X

P dt,j i

k X

+ 2ηl2 E

j=0

P ξit,j

(10)

j=0

This step does not use any independence across local steps; it is purely deterministic. Step 3: Bounding the drift part (deterministic terms). m−1 X

≤m

m−1 X

j=0

Let m := k + 1. For any vectors a0 , . . . , am−1 , ∥aj ∥2 .

j=0

Applying this with aj = P dt,j i and using ∥P v∥ ≤ ∥v∥, we get k X

P dt,j i

≤m

j=0

k X

E dt,j i

(11)

j=0

We have t,j t t t t dt,j i = ∇Fi (yi ) − ∇Fi (x ) + ∇Fi (x ) − ci + c . Using ∥u + v∥2 ≤ 2∥u∥2 + 2∥v∥2 and L-smoothness (Assumption 1), t,j 2 t ∥dt,j i ∥ ≤ 2 ∇Fi (yi ) − ∇Fi (x )

≤ 2L2 yit,j − xt

+ 2 ∇Fi (xt ) − cti + ct

+ 2Gi,t ,

where Gi,t := ∇Fi (xt ) − cti + ct

t,j By the reconstruction yit,j = P ⊤ yi,proj + xtres and xt = P ⊤ xtproj + xtres , we have yit,j − xt = P ⊤ ∆t,j i and hence t,j t,j 2 t 2 ∥yi − x ∥ ≤ ∥∆i ∥ . Therefore 2

E dt,j i

2 ≤ 2L2 E∥∆t,j i ∥ + 2Gi,t .

Substituting into (11), E

k X

P dt,j i

j=0

≤ 2L2 m

k X

2 2 E∥∆t,j i ∥ + 2m Gi,t .

(12)

j=0

Step 4: Bounding the noise part via conditional independence. Let Ht denote the sigma-algebra generated by all randomness up to the beginning of round t, including the projector P , the global iterate xt , and the control variates, but **before** sampling the stochastic gradients {ξit,j }K−1 j=0 for that round. By Assumption (A2), conditioned on Ht the noise variables {ξit,j }kj=0 are independent, have zero mean and satisfy E[∥ξit,j ∥2 | Ht ] ≤ σ 2 . We first compute the conditional second moment and then take a full expectation: E

k X

P ξit,j

k h h X ii 2 =E E P ξit,j Ht

j=0

(13)

j=0 k hX =E E ∥P ξit,j ∥2 | Ht + 2 j=0

i E ⟨P ξit,j , P ξit,ℓ ⟩ | Ht

(14)

0≤j<ℓ≤k

k hX i =E E ∥P ξit,j ∥2 | Ht

(15)

j=0 k hX i ≤E E ∥ξit,j ∥2 | Ht ≤ m σ 2 , j=0

where we used ∥P v∥ ≤ ∥v∥ and the single-step variance bound in Assumption (A2). 10

(16)

2 Step 5: Recursion for ek := E∥∆t,k i ∥ .

Combining the drift and noise bounds with (10), we obtain

k X 2 2 2 ∥ ≤ 2η 2L m ek+1 := E∥∆t,k+1 ej + 2m2 Gi,t + mσ 2 l i

(17)

j=0

= 4ηl2 L2 m

k X

ej + 4ηl2 m2 Gi,t + 2ηl2 m σ 2 .

(18)

j=0

For k ≤ K − 1 we have m = k + 1 ≤ K, so m ≤ K and m2 ≤ K 2 ; hence ek+1 ≤ α

k X

ej + C,

(19)

j=0

with α := 4ηl2 L2 K, Let Sk :=

C := 4ηl2 K 2 Gi,t + 2ηl2 K σ 2 .

j=0 ej . Then S0 = e0 = 0 and

Sk+1 ≤ (1 + α)Sk + C. By induction, Sk ≤

C C (1 + α)k − 1 ≤ (1 + α)k . α α

Using this in ek+1 ≤ αSk + C gives ek+1 ≤ C(1 + α)k . Under the stepsize condition ηl ≤ 1/(2KL) we have α = 4ηl2 L2 K ≤ 1/K, so (1 + α)k ≤ (1 + 1/K)K ≤ e < 3 for k ≤ K. Thus, ek+1 ≤ 3C = 12ηl2 K 2 Gi,t + 6ηl2 Kσ 2 , 0 ≤ k ≤ K − 1. (20) Step 6: Summing over local steps. Ei,t =

K−1 X

ek =

k=0

K X

Finally, ek ≤ K max ek ≤ K 12ηl2 K 2 Gi,t + 6ηl2 Kσ 2 ,

k=1

1≤k≤K

which yields Ei,t ≤ 12K 3 ηl2 ∇Fi (xt ) − cti + ct

+ 6K 2 ηl2 σ 2 .

This proves the claimed bound (8). Throughout the proof we used Assumption (A1) for smoothness and Assumption (A2) both for single-step variance control and, crucially, for the independence of the local noises {ξit,k }k across steps within the same round, which is what allows the noise term to scale as O(K 2 ηl2 σ 2 ) instead of O(K 3 ηl2 σ 2 ).

4.5

One-round progress

Lemma 3 (One-Round Progress). Suppose Assumptions 1–3 hold, and let the effective global step size be η̃ := ηg ηl K. Define the average projected drift N K−1

t Edrift :=

1 XX 2 t,k E yi,proj − xtproj , N K i=1 k=0

1 which is controlled by Lemma 2. If η̃ ≤ 4L , then the one-round progress satisfies

ρ η̃ Lη̃ 2 2 3L2 η̃ t 2 E F (xt+1 ) ≤ F (xt ) − ∇F (xt ) + σ + Edrift . 4 2N K 2

(21)

Here the expectation is taken over the randomness of the projector Pt , the sampled clients, and all stochastic gradients in round t. 11

Proof. We first rewrite the global update in a form that makes the projected direction explicit. From Algorithm 1, the projected local update on client i is t,k+1 t,k yi,proj = yi,proj − ηl Pt git,k − cti,proj + ctproj , t,k t,0 where git,k = ∇Fi (yit,k ; ξit,k ) and yit,k = Pt⊤ yi,proj +xtres . Unrolling over k = 0, . . . , K−1 and using yi,proj = xtproj gives K−1 X t,K xtproj − yi,proj = ηl Pt git,k − cti,proj + ctproj . (22) k=0

The projected global aggregation step reads t xt+1 proj = xproj − ηg

t,K N ηl K X xtproj − yi,proj . N i=1 ηl K

Substituting (22) and simplifying yields N

t xt+1 proj = xproj − ηg

1 X t t,K xproj − yi,proj N i=1

= xtproj − ηg ηl

(23)

N K−1 1 XX Pt git,k − cti,proj + ctproj . N i=1

(24)

k=0

Define the averaged projected direction N K−1

vt :=

1 XX Pt git,k − cti,proj + ctproj ∈ Rr . N K i=1 k=0

Then t xt+1 proj = xproj − η̃vt ,

η̃ := ηg ηl K,

and the backfill step gives t t ⊤ xt+1 = Pt⊤ xt+1 proj + xres = x − η̃Pt vt .

Using

(25)

t t i=1 ci,proj = N cproj , the control variates cancel in vt and we obtain N K−1

vt =

1 XX Pt git,k . N K i=1 k=0

Using the orthogonal projector Pt onto the random r-dimensional subspace in Rd (see Section 4.2), we have Pt⊤ vt = Pt

N K−1 1 X X t,k gi . N K i=1

(26)

k=0

Decomposition into mean gradient and noise. where by Assumption 2 E[ξit,k | yit,k ] = 0,

Write the stochastic gradient as git,k = ∇Fi (yit,k ) + ξit,k , E ∥ξit,k ∥2 | yit,k ≤ σ 2 ,

and the random noises {ξit,k }i,k are independent across different clients i within the same round t, as specified in Assumption 2. Define the averaged full gradients and noise in round t: N K−1

ḡt :=

1 XX ∇Fi (yit,k ), N K i=1

N K−1

ξ¯t :=

k=0

1 X X t,k ξi . N K i=1 k=0

Then (26) becomes Pt⊤ vt = Pt (ḡt + ξ¯t ),

(27)

xt+1 = xt − η̃ Pt (ḡt + ξ¯t ).

(28)

so the full-space update (25) is

Smoothness-based descent inequality.

By L-smoothness of F (Assumption 1),

L t+1 2 x − xt 2 Lη̃ 2 2 = F (xt ) − η̃ ∇F (xt ), Pt (ḡt + ξ¯t ) + Pt (ḡt + ξ¯t ) . 2

F (xt+1 ) ≤ F (xt ) + ∇F (xt ), xt+1 − xt +

(29) (30)

Denote Gt := ∇F (xt ). We decompose the right-hand side into the linear term T1 := −η̃ Gt , Pt (ḡt + ξ¯t ) and the quadratic term T2 :=

Lη̃ 2 2 Pt (ḡt + ξ¯t ) . 2

We now bound E[T1 ] and E[T2 ]. Bounding the linear term T1 . We first note that xt is fixed before sampling Pt and the stochastic gradients in round t. Using (28), T1 = −η̃ Gt , Pt ḡt − η̃ Gt , Pt ξ¯t .

(31)

Conditioned on the sigma-fields that encode the past iterates, each local noise has zero conditional mean. More precisely, let Ft be the sigma-algebra generated by all randomness up to the beginning of round t (so that xt , Gt and Pt are Ft -measurable), and for each client i and local step k let Ft,i,k be the sigma-algebra generated by Ft together with all local iterates up to step k on all clients. Then yit,k is Ft,i,k -measurable and Assumption (A2) implies E[ξit,k | Ft,i,k ] = 0. Using the tower property, we get E ⟨Gt , Pt ξ¯t ⟩ =

N K−1 1 XX E ⟨Gt , Pt ξit,k ⟩ N K i=1 k=0

N K−1 X X

1 N K i=1

k=0

N K−1 X X

1 N K i=1

h i E E ⟨Gt , Pt ξit,k ⟩ | Ft,i,k h i E Gt , Pt E[ξit,k | Ft,i,k ] = 0.

k=0

For the gradient part, decompose ḡt = Gt + (ḡt − Gt ). Then E[T1 ] = −η̃E ⟨Gt , Pt Gt ⟩ − η̃E ⟨Gt , Pt (ḡt − Gt )⟩ .

(32)

Using Assumption (A3), E[Pt ] = ρId with ρ = r/d, so E ⟨Gt , Pt Gt ⟩ = ⟨Gt , E[Pt ]Gt ⟩ = ρ ∥Gt ∥2 . For the second term, we use Cauchy–Schwarz and Young’s inequality together with ∥Pt z∥ ≤ ∥z∥ and E∥Pt Gt ∥2 = ρ ∥Gt ∥2 : −E ⟨Gt , Pt (ḡt − Gt )⟩ = −E ⟨Pt Gt , ḡt − Gt ⟩ (33) 1 1 ≤ E∥Pt Gt ∥2 + E∥ḡt − Gt ∥2 (34) 2 2 ρ 1 = ∥Gt ∥2 + E∥ḡt − Gt ∥2 . (35) 2 2

Combining the two pieces, we obtain E[T1 ] ≤ −

ρη̃ η̃ ∥Gt ∥2 + E∥ḡt − Gt ∥2 . 2 2

(36)

We now bound the bias term E∥ḡt − Gt ∥2 by the drift. Observe that ḡt − Gt =

N K−1 1 XX ∇Fi (yit,k ) − ∇Fi (xt ) . N K i=1 k=0

By Jensen’s inequality and L-smoothness, N K−1

1 XX 2 ∇Fi (yit,k ) − ∇Fi (xt ) ∥ḡt − Gt ∥ ≤ N K i=1 2

(37)

k=0

N K−1

≤

L2 X X t,k 2 yi − x t . N K i=1

(38)

k=0

t,k t,k The reconstruction formulas yit,k = Pt⊤ yi,proj +xtres and xt = Pt⊤ xtproj +xtres imply yit,k −xt = Pt⊤ (yi,proj −xtproj ), so 2 2 t,k t,k ∥yit,k − xt ∥2 = Pt⊤ (yi,proj − xtproj ) ≤ yi,proj − xtproj . t Taking expectations and using the definition of Edrift , we get t E∥ḡt − Gt ∥2 ≤ L2 Edrift .

(39)

Substituting into (36), E[T1 ] ≤ − Bounding the quadratic term T2 .

ρη̃ η̃L2 t ∥Gt ∥2 + E . 2 2 drift

(40)

We have Lη̃ 2 2 E Pt (ḡt + ξ¯t ) 2 Lη̃ 2 Lη̃ 2 2 2 = E Pt ḡt + E Pt ξ¯t , 2 2

E[T2 ] =

(41) (42)

where we used that the mixed term E⟨Pt ḡt , Pt ξ¯t ⟩ vanishes. Indeed, expanding ξ¯t and using the same filtration Ft,i,k as above, E ⟨Pt ḡt , Pt ξ¯t ⟩ =

N K−1 1 XX E ⟨Pt ḡt , Pt ξit,k ⟩ N K i=1 k=0

N K−1 X X

1 N K i=1

k=0

N K−1 X X

1 N K i=1

h i E E ⟨Pt ḡt , Pt ξit,k ⟩ | Ft,i,k h i E Pt ḡt , Pt E[ξit,k | Ft,i,k ] = 0.

k=0

Gradient part. Using ḡt = Gt + (ḡt − Gt ) and the inequality ∥a + b∥2 ≤ 2∥a∥2 + 2∥b∥2 , Pt ḡt

≤ 2 Pt Gt

+ 2 Pt (ḡt − Gt ) + 2 ḡt − Gt

(43) (44)

t Taking expectations and using E∥Pt Gt ∥2 = ρ ∥Gt ∥2 and E∥ḡt − Gt ∥2 ≤ L2 Edrift , we obtain

E Pt ḡt

t ≤ 2ρ ∥Gt ∥2 + 2L2 Edrift .

(45)

Noise part (using independence across clients and local steps). Conditioned on Pt , let Σξ̄|Pt denote the (conditional) covariance matrix of ξ¯t . Then E ∥Pt ξ¯t ∥2 | Pt = E ξ¯t⊤ Pt⊤ Pt ξ¯t | Pt (46) = tr Pt Σξ̄|Pt . (47) Using ∥Pt ∥2 ≤ 1 and the fact that Σξ̄|Pt is positive semi-definite, tr(Pt Σξ̄|Pt ) ≤ ∥Pt ∥2 tr(Σξ̄|Pt ) ≤ tr(Σξ̄|Pt ) = E ∥ξ¯t ∥2 | Pt . Taking expectations over Pt gives the general bound 2 2 E Pt ξ¯t ≤ E ξ¯t ,

(48)

which does not require any independence between Pt and the gradient noise. To bound E∥ξ¯t ∥2 , we now use the conditional independence across both clients and local steps in Assumption (A2). Let Ht denote the sigma-algebra generated by the past iterates up to the beginning of round t. Then, conditioned on Ht , the family {ξit,k }i,k is independent, mean-zero, and satisfies E[∥ξit,k ∥2 | Ht ] ≤ σ 2 . By definition, N K−1

1 X X t,k 2 ξi N K i=1

(49)

N K−1 h X X t,k 2 i 1 ξi E 2 2 N K i=1

(50)

N K−1 i h X X t,k 2 1 ξ H . E E t i N 2K 2 i=1

(51)

2 E ξ¯t = E

k=0

Inside the inner expectation, conditional independence and zero mean yield cancellation of all cross terms, so E ξ¯t

N K−1 hX X t,k i 1 2 E E ∥ξ ∥ H t i N 2 K 2 i=1 k=0 h i 1 σ2 . ≤ 2 2 E N K σ2 = N K NK

(52) (53)

Combining with (48), we obtain σ2 2 . E Pt ξ¯t ≤ NK

(54)

Putting (45) and (54) together, we get Lη̃ 2 E[T2 ] ≤ 2

σ2 2ρ ∥Gt ∥ + 2L NK 2 Lη̃ t = Lη̃ 2 ρ ∥Gt ∥2 + L3 η̃ 2 Edrift + σ2 . 2N K

Combining T1 and T2 .

t Edrift +

(55) (56)

Combining (40) and (56), we obtain ρη̃ E F (xt+1 ) − F (xt ) ≤ − ∥Gt ∥2 + Lη̃ 2 ρ ∥Gt ∥2 2 η̃L2 Lη̃ 2 2 t + L3 η̃ 2 Edrift + σ . + 2 2N K

(57) (58)

The coefficient of ∥Gt ∥2 can be written as −ρη̃

1 2

− Lη̃ .

1 , we have 21 − Lη̃ ≥ 14 , hence Under the step-size condition η̃ ≤ 4L

−

ρη̃ ρη̃ + Lη̃ 2 ρ ≤ − . 2 4

(59)

For the drift term, note that η̃L2 3η̃L2 3η̃L2 η̃L2 η̃L2 + L3 η̃ 2 = (1 + 2Lη̃) ≤ (1 + 2 · 14 ) = ≤ . 2 2 2 4 2 Therefore, ρη̃ Lη̃ 2 2 3L2 η̃ t E F (xt+1 ) − F (xt ) ≤ − ∥∇F (xt )∥2 + σ + Edrift . 4 2N K 2 Rearranging yields the claimed bound (21). This completes the proof.

4.6

Control variate contraction

Lemma 4 (Control Variate Contraction). Suppose Assumptions 1–3 hold, and that the local stepsize satisfies ηl ≤ 1/(2KL) so that Lemma 2 applies. Define the control-variate error N

Ct :=

1 X 2 E cti − ∇Fi (xt ) , N i=1

and the averaged projected drift as in Lemma 3, N K−1

t Edrift :=

1 XX 2 t,k E yi,proj − xtproj . N K i=1 k=0

Let η̃ := ηg ηl K be the effective global stepsize. Then, for every round t, the control variates satisfy ρ σ2 t Ct+1 ≤ 1 − Ct + 3 + 3L2 Edrift 2 K 12L4 2 t 6L2 η̃ 2 σ 2 12L2 2 2 η̃ ∇F (xt ) + + η̃ Edrift . + ρ ρ NK ρ

(60)

In particular, Ct contracts by a factor 1 − ρ/2 up to perturbation terms that depend on the stochastic variance t σ 2 , the drift Edrift , and the model movement η̃ 2 ∥∇F (xt )∥2 . Proof. Fix an outer round t and a client i. Write vi∗ := ∇Fi (xt ) for the reference gradient at the beginning of round t. We first track the evolution of ct+1 with respect to the frozen reference vi∗ , then relate vi∗ to the new i t+1 reference ∇Fi (x ). From Algorithm 1 and the projector Pt

Step 1: Subspace decomposition of the control update. defined in Section 4.2, the control update is ct+1 = (I − Pt )cti + Pt i

1 K−1 X K

git,k ,

(61)

k=0

where git,k = ∇Fi (yit,k ; ξit,k ) is the stochastic gradient at local step k on client i. Using vi∗ = (I − Pt )vi∗ + Pt vi∗ , we obtain ct+1 − vi∗ = (I − Pt )(cti − vi∗ ) + Pt i

1 K−1 X K

k=0

git,k − vi∗ .

(62)

Define Ai := (I − Pt )(cti − vi∗ ),

Bi := Pt (ḡit − vi∗ ),

PK−1 t,k 1 where ḡit := K k=0 gi . For a fixed projector Pt , we have Ai ∈ ker(Pt ) and Bi ∈ range(Pt ), and therefore Ai ⊥ Bi . Hence 2 ct+1 − vi∗ = ∥Ai ∥2 + ∥Bi ∥2 . (63) i Step 2: Contraction in the orthogonal complement via (A3). Conditioning on the history up to the beginning of round t, both cti and vi∗ are fixed before sampling Pt . Taking expectation over Pt and using Assumption 3, EPt ∥Ai ∥2 = EPt (cti − vi∗ )⊤ (I − Pt )(cti − vi∗ ) (64) t t ∗ ⊤ ∗ = (ci − vi ) I − EPt [Pt ] (ci − vi ) (65) = (1 − ρ) cti − vi∗

(66)

Taking full expectation we obtain E∥Ai ∥2 = (1 − ρ) E cti − ∇Fi (xt )

(67)

Step 3: Bounding the projected gradient error Bi using (A1) and (A2). We decompose the error in ḡit into a stochastic noise part and a drift part. Write git,k = ∇Fi (yit,k ) + ξit,k , where, by Assumption 2, E[ξit,k | yit,k ] = 0,

E ∥ξit,k ∥2 | yit,k ≤ σ 2 ,

and the noises {ξit,k } are independent across different local steps k on the same client and across different clients. Define K−1 K−1 1 X t,k 1 X ξ¯i := ξi , ∇Fi (yit,k ) − ∇Fi (xt ) . δ̄i := K K k=0

k=0

ḡit − vi∗ = ξ¯i + δ̄i ,

Bi = Pt (ξ¯i + δ̄i ).

Then Using ∥u + v∥2 ≤ 2∥u∥2 + 2∥v∥2 and ∥Pt z∥ ≤ ∥z∥, we get E∥Bi ∥2 ≤ 2 E∥Pt ξ¯i ∥2 + 2 E∥Pt δ̄i ∥2 ≤ 2 E∥Pt ξ¯i ∥2 + 2 E∥δ̄i ∥2 . Noise part. Condition on Pt and let Σξ̄i |Pt be the conditional covariance of ξ¯i . Then E ∥Pt ξ¯i ∥2 | Pt = tr Pt Σξ̄i |Pt ≤ ∥Pt ∥2 tr Σξ̄i |Pt ≤ tr Σξ̄i |Pt = E ∥ξ¯i ∥2 | Pt .

(68)

(69) (70) (71)

Taking expectation over Pt yields the general bound E∥Pt ξ¯i ∥2 ≤ E∥ξ¯i ∥2 .

(72)

By independence of the local noises across k (Assumption 2), K−1

E∥ξ¯i ∥2 = E

1 X t,k 2 ξi K

(73)

k=0

K−1

1 X σ2 ≤ 2 E∥ξit,k ∥2 ≤ . K K k=0

(74)

Combining with (72), σ2 . K Drift part. By Jensen’s inequality and L-smoothness (Assumption 1), E∥Pt ξ¯i ∥2 ≤

δ̄i

(75)

K−1 2 1 X ∇Fi (yit,k ) − ∇Fi (xt ) K

(76)

k=0

K−1

≤

1 X 2 ∇Fi (yit,k ) − ∇Fi (xt ) K

(77)

k=0

K−1

L2 X t,k 2 ≤ yi − xt . K

(78)

k=0

t,k Using the reconstruction formulas yit,k = Pt⊤ yi,proj + xtres and xt = Pt⊤ xtproj + xtres we get yit,k − xt = t,k t,k t,k Pt⊤ (yi,proj − xtproj ) and hence ∥yi − xt ∥2 ≤ ∥yi,proj − xtproj ∥2 . Thus, K−1

E δ̄i

L2 L2 X 2 t,k − xtproj = Ei,t , E yi,proj K K

≤

(79)

k=0

where Ei,t is as in Lemma 2 and obeys the refined drift bound Ei,t ≤ 6K 2 ηl2 σ 2 + 12K 3 ηl2 ∇Fi (xt ) − cti + ct Combining (68), (75) and the drift bound gives E∥Bi ∥2 ≤ 2

σ2 2L2 + Ei,t . K K

(80)

Step 4: Recursion with respect to the frozen reference vi∗ (using Lemma 2). Taking expectations in (63) and using (67) and (80), we obtain 2L2 σ2 + Ei,t . K K This step relies on Assumptions (A1)–(A3) and the drift quantity Ei,t controlled in Lemma 2. E ct+1 − vi∗ i

≤ (1 − ρ) E cti − vi∗

(81)

Step 5: Changing the reference to ∇Fi (xt+1 ). We now relate vi∗ = ∇Fi (xt ) to the next-round reference ∇Fi (xt+1 ). By Young’s inequality, for any β > 0, 1 ∗ 2 t+1 t+1 2 ∗ 2 vi − ∇Fi (xt+1 ) . (82) ct+1 − ∇F (x ) ≤ (1 + β) c − v + 1 + i i i i β Choose β = ρ/2. Using L-smoothness, vi∗ − ∇Fi (xt+1 )

= ∇Fi (xt ) − ∇Fi (xt+1 )

≤ L2 xt+1 − xt

Moreover, for ρ ∈ (0, 1] we have (1 + ρ/2)(1 − ρ) = 1 − ρ/2 − ρ2 /2 ≤ 1 − ρ/2 and 1 + ρ/2 ≤ 3/2, 1 + 1/β = 1 + 2/ρ ≤ 3/ρ. Therefore, inserting (81), h i σ2 2L2 2 2 E ct+1 − ∇Fi (xt+1 ) ≤ (1 + ρ2 ) (1 − ρ) E cti − ∇Fi (xt ) + 2 + Ei,t (83) i K K 2 2 2 + 1+ L E xt+1 − xt (84) ρ ρ 3L2 σ2 2 ≤ 1− E cti − ∇Fi (xt ) + 3 + Ei,t (85) 2 K K 3L2 2 E xt+1 − xt . (86) + ρ t Averaging over i and recalling the definitions of Ct and Edrift , 2 ρ σ 3L2 2 t Ct+1 ≤ 1 − Ct + 3 + 3L2 Edrift + E xt+1 − xt . 2 K ρ

(87)

Step 6: Bounding the model movement using Lemma 3. We now bound E∥xt+1 − xt ∥2 using the global update structure explicitly derived in the proof of Lemma 3. In particular, as shown there (cf. equation (28) in that proof), the full-space update can be written as xt+1 = xt − η̃ Pt (ḡt + ξ¯t ), where

N K−1

ḡt :=

(88) N K−1

1 XX ∇Fi (yit,k ), N K i=1

ξ¯t :=

k=0

1 X X t,k ξi . N K i=1 k=0

Using ∥Pt z∥ ≤ ∥z∥, we obtain E xt+1 − xt

= η̃ 2 E Pt (ḡt + ξ¯t ) ≤ η̃ 2 E ḡt + ξ¯t

(89)

(90)

≤ 2η̃ E∥ḡt ∥ + 2η̃ E∥ξ¯t ∥2 . 2

(91)

Mean-gradient part. Let Gt := ∇F (xt ). Then E∥ḡt ∥2 = E Gt + (ḡt − Gt ) ≤ 2 Gt

+ 2 E ḡt − Gt

(92) 2

(93)

t , we obtain Exactly as in Lemma 3, using Jensen’s inequality, L-smoothness, and the definition of Edrift

E ḡt − Gt

t ≤ L2 Edrift .

(94)

Therefore, E∥ḡt ∥2 ≤ 2 ∇F (xt )

t + 2L2 Edrift .

(95)

Noise part (global averaging). By Assumption (A2) and independence of stochastic gradients across clients and local steps, N K−1

E ξ¯t

1 X X t,k 2 ξi N K i=1

(96)

k=0

≤

N K−1 σ2 1 XX t,k 2 ≤ ξ E . i N 2 K 2 i=1 NK

(97)

k=0

Combining the two parts, E xt+1 − xt

σ2 2 t ≤ 2η̃ 2 2 ∇F (xt ) + 2L2 Edrift + 2η̃ 2 NK 2 σ 2 t = 4η̃ 2 ∇F (xt ) + 4L2 η̃ 2 Edrift + 2η̃ 2 . NK

(98) (99)

This bound again relies only on Assumptions (A1) and (A2) and the update representation shared with Lemma 3; it does not assume any independence between Pt and the noise. Step 7: Plugging the movement bound into the control recursion. Substituting the last inequality into (87), we obtain ρ σ2 t Ct+1 ≤ 1 − Ct + 3 + 3L2 Edrift (100) 2 K σ2 3L2 2 2 t + 4η̃ ∇F (xt ) + 4L2 η̃ 2 Edrift + 2η̃ 2 . (101) ρ NK Rearranging terms yields exactly the claimed recursion (60). This proves that Ct contracts by a factor 1 − ρ/2 t up to variance and drift terms whose dependence on σ 2 , Edrift , and η̃ 2 ∥∇F (xt )∥2 matches the structure used in the one-round progress analysis. 19

4.7

Main convergence theorem

Theorem 1 (Global Convergence of Subspace SCAFFOLD). Suppose Assumptions 1, 2, and 3 hold. Let ρ = r/d and ∆F = F (x0 ) − F ∗ . Choose the local learning rate ηl and the effective global stepsize η̃ = ηg ηl K such that 1 1 1. ηl ≤ 2KL and η̃ ≤ 4L ;

2. 72K 3 L2 ηl2 ≤ ρ8 ; 3.

3456L4 η̃ 2 K 3 ηl2 ≤ 1; ρ3 2 2

ρ ). 4. 4Lρ2η̃ ≤ 1 (equivalently, η̃ ≤ 2L

Then the iterates generated by Algorithm 1 satisfy min E ∇F (xt )

0≤t<T

≤

8(∆F + C0 ) 4Lσ 2 η̃ 8000L2 Kηl2 σ 2 + + . ρ η̃ T ρ NK ρ

(102)

In particular, the bound has the standard nonconvex stochastic-gradient structure O (ρη̃T )−1 + (ρ−1 N −1 K −1 )η̃ + (ρ−1 Kηl2 ) . Proof. Define the Lyapunov function Φt := F (xt ) − F ∗ + M Ct , where M > 0 will be specified later. We analyze E[Φt+1 ] − Φt . All expectations below are taken over the randomness of the projector Pt , the stochastic gradients, and the client sampling in round t. Step 1: One-round progress for F (Lemma 3).

By Lemma 3, if η̃ ≤ 1/(4L), then

ρ η̃ Lη̃ 2 2 3L2 η̃ t 2 E F (xt+1 ) − F (xt ) ≤ − ∇F (xt ) + σ + Edrift . 4 2N K 2

(103)

Compared with earlier one-round bounds, the variance term no longer carries a factor ρ; this is the more conservative but dimensionally sharp version provided by the updated Lemma 3. Step 2: Control variate contraction (Lemma 4).

By Lemma 4, for every round t we have

ρ σ2 t Ct+1 ≤ 1 − + 3L2 Edrift Ct + 3 2 K 6L2 η̃ 2 σ 2 12L2 2 12L4 2 t 2 + η̃ ∇F (xt ) + + η̃ Edrift . ρ ρ NK ρ

(104)

This already incorporates the refined drift analysis and the updated dependence on η̃ and ρ from Lemma 4. Step 3: Drift of the Lyapunov function.

Multiplying (104) by M and adding it to (103) yields

ρ η̃ Lη̃ 2 2 3L2 η̃ t 2 σ + Edrift ∇F (xt ) + 4 2N K 2 h ρ σ2 t + M − Ct + 3 + 3L2 Edrift 2 K 12L2 2 6L2 η̃ 2 σ 2 12L4 2 t i 2 + η̃ ∇F (xt ) + + η̃ Edrift . ρ ρ NK ρ

E[Φt+1 ] − Φt ≤ −

t Grouping terms according to ∥∇F (xt )∥2 , Ct , and Edrift , we obtain

ρ η̃ 12M L2 Mρ 2 E[Φt+1 ] − Φt ≤ − − η̃ 2 ∇F (xt ) − Ct 4 ρ 2 20

(105)

t + Bdrift Edrift + Cnoise ,

(106)

where 3L2 η̃ 12L4 2 + M 3L2 + η̃ , 2 ρ σ2 Lη̃ 2 2 6L2 η̃ 2 σ 2 Cnoise := . σ +M 3 + 2N K K ρ NK

(107)

Bdrift :=

(108)

1 Step 4: Using the refined drift bound (Lemma 2). Lemma 2 states that, if ηl ≤ 2KL , then for every client i and round t,

Ei,t :=

K−1 X

t,k E yi,proj − xtproj

≤ 6K 2 ηl2 σ 2 + 12K 3 ηl2 ∇Fi (xt ) − cti + ct

k=0

Recalling N

t Edrift =

1 X Ei,t , N K i=1

we obtain t Edrift ≤ 6Kηl2 σ 2 + 12K 2 ηl2 Gt ,

where

(109)

1 X 2 Gt := E ∇Fi (xt ) − cti + ct . N i=1 As in the lemma proof, we can relate Gt to the control-variate error Ct and the full gradient via Gt ≤ 2 Ct + 2 ∇F (xt )

(110)

Combining (109) and (110) yields t Edrift ≤ 6Kηl2 σ 2 + 24K 2 ηl2 Ct + 24K 2 ηl2 ∇F (xt )

(111)

t Compared with the old drift bound, the noise term now scales as O(Kηl2 σ 2 ) in Edrift (instead of O(K 2 ηl2 σ 2 )), t 2 2 2 while the dependence on Ct and ∥∇F (x )∥ remains of order K ηl . Substituting (111) into (106), and regrouping the terms proportional to ∥∇F (xt )∥2 and Ct , we obtain

E[Φt+1 ] − Φt ≤ −ag ∇F (xt )

− a c Ct + N t ,

(112)

where ag :=

ρ η̃ 12M L2 2 − η̃ − 24Bdrift K 2 ηl2 , 4 ρ

(113)

ac :=

Mρ − 24Bdrift K 2 ηl2 , 2

(114)

Nt := Cnoise + 6Bdrift Kηl2 σ 2 .

(115)

and the accumulated noise term is

Step 5: Choice of M and verification of descent. We now choose M and enforce conditions (2)–(4) in the theorem statement so that both ag and ac are uniformly positive. Following the same design as in the previous version, set 72L2 η̃ K 3 ηl2 M := . (116) ρ

Using this definition together with (107) and (111), one checks that Conditions (2)–(4) imply ag ≥

ρ η̃ , 8

ac ≥

Mρ . 4

(117)

(Here we only use that K 3 L2 ηl2 and L2 η̃ 2 /ρ2 are sufficiently small, as encoded in Conditions (2)–(4).) Therefore, from (112), ρ η̃ Mρ 2 E[Φt+1 ] − Φt ≤ − ∇F (xt ) − Ct + N t . (118) 8 4 Step 6: Bounding the total noise term Nt . We now bound Nt in terms of σ 2 and the algorithmic parameters. From (107), (108), and (115), using the choice (116) and Conditions (3)–(4) to dominate higher-order terms in η̃, we obtain Nt ≤

Lη̃ 2 2 σ + c2 L2 η̃ Kηl2 σ 2 , {z } 2N K | {z } |

(119)

=: N2

=: N1

for some universal numerical constant c2 > 0; a conservative choice is c2 = 1000, which is compatible with the refined drift bound (111). Here N1 collects the direct variance contributions from Lemma 3 and Lemma 4, while N2 collects those induced via the (drift-dependent) terms in both lemmas. Note that the improved drift noise 6Kηl2 σ 2 in (111) indeed yields N2 of order L2 η̃Kηl2 σ 2 . Step 7: Telescoping sum and the final convergence bound. Summing (118) over t = 0, 1, . . . , T − 1 and using Φt ≥ 0 gives T −1 T −1 X ρ η̃ X 2 Nt E ∇F (xt ) ≤ Φ0 − E[ΦT ] + 8 t=0 t=0

≤ Φ0 + T N̄ , where N̄ := max0≤t<T Nt . From the definition of Φt and the choice of M , we have Φ0 ≤ ∆F + C0 up to a factor absorbed into the constant 8. PT −1 Dividing both sides by T and using min0≤t<T at ≤ T1 t=0 at with at = E∥∇F (xt )∥2 , we obtain min E ∇F (xt )

0≤t<T

8 8(∆F + C0 ) + N̄ ρ η̃ T ρ η̃ 8(∆F + C0 ) 8 Lη̃ 2 2 2 2 2 ≤ + σ + c2 L η̃ Kηl σ ρ η̃ T ρ η̃ 2N K 8(∆F + C0 ) 4Lσ 2 η̃ 8000L2 Kηl2 σ 2 ≤ + + , ρ η̃ T ρ NK ρ ≤

where in the last step we used c2 ≤ 1000 and absorbed numerical factors into the constant 8000. This is precisely the bound (102) with an explicit universal constant C = 8000 in front of the L2 Kηl2 σ 2 /ρ term. Independence structure. Finally, we emphasise that Lemmas 2, 3, and 4, and hence the entire argument above, never assume any independence between the random projector Pt and the stochastic gradients or noise. The only probabilistic assumptions used are Assumptions 1–3 and the standard independent mini-batch sampling across clients and local steps. The price for removing any projector–noise independence assumption is the somewhat more conservative variance term in (102); the improved drift bound in Lemma 2 is fully reflected in the L2 η̃Kηl2 σ 2 scaling of N2 and the final L2 Kηl2 σ 2 /ρ term.

4.8

Harmonic-step-size corollary

Corollary 1 (Non-asymptotic rate with harmonic stepsize and split local/global rates). Suppose Assumptions 1, 2, and 3 hold and let ρ = r/d ∈ (0, 1] and ∆F = F (x0 ) − F ∗ . Consider Algorithm 1 with constant 22

stepsizes. Let C0 be the initial control-variate error appearing in Theorem 1. Fix an absolute constant C∗ := 100. We choose the stepsizes as follows: 1. The local stepsize ηl is set to r n 1 ρ , , ηl := min 2KL 864 L2 K 3

∆F + C 0 o . C∗ LKσ 2 T ρ

(120)

2. The global stepsize is set to ηg := ηlη̃K , where the effective stepsize η̃ is defined by the harmonic rule −1 η̃ := η̃0−1 + η̃1−1 ,

(121)

using the stability threshold η̃0 and the variance-reduction threshold η̃1 : r n 1 ρ o 2N K(∆F + C0 ) η̃0 := min , η̃1 := , . 4L 2L Lσ 2 T

(122)

Under these choices, the iterates of Algorithm 1 satisfy t

min E ∇F (x )

0≤t<T

√ p 112 L ∆F + C0 8 2 σ √ ≤ + L(∆F + C0 ). ρ2 T ρ N KT

(123)

In particular, the convergence rate is 2

min E ∇F (x )

0≤t<T

L σ = O 2 F (x0 ) − F ∗ + C0 + √ ρ T ρ N KT

! L F (x0 ) − F ∗ + C0

(124)

Proof. Step 1: Starting from Theorem 1. By Theorem 1, under Assumptions 1–3 and the step-size conditions (1)–(4), the iterates of Algorithm 1 satisfy min E ∇F (xt )

0≤t<T

≤

8(∆F + C0 ) 4Lσ 2 η̃ 8000L2 Kηl2 σ 2 + + . ρ η̃ T ρ NK ρ

(125)

We define the coefficients A :=

8(∆F + C0 ) , ρ

B :=

4Lσ 2 , ρN K

D :=

8000L2 Kσ 2 . ρ

(126)

Then the bound (125) can be rewritten as min E ∇F (xt )

0≤t<T

≤ Φ(η̃, ηl ) :=

A + B η̃ + D ηl2 . η̃T

We observe that the threshold η̃1 defined in (122) satisfies s r r 2N K(∆F + C0 ) 8(∆F + C0 )/ρ A η̃1 = = = . Lσ 2 T (4Lσ 2 /(ρN K)) · T BT

(127)

(128)

Step 2: Verifying the step-size constraints. We check that the choices of ηl and η̃ are compatible with Conditions (1)–(4) of Theorem 1. 1 (i) Condition (1). By construction in (120), ηl ≤ 2KL , so the first part holds. Moreover, by (121), η̃ ≤ η̃0 . Since η̃0 ≤ 1/(4L) by (122), we have η̃ ≤ 1/(4L). (ii) Condition (4). Since η̃ ≤ η̃0 ≤ ρ/(2L), we have 4L2 η̃ 2 /ρ2 ≤ 1, so Condition (4) holds. (iii) Condition (2). By (120), we have ηl2 ≤

ρ , 864L2 K 3 23

which implies 72K 3 L2 ηl2 ≤ ρ/12 ≤ ρ/8. Thus Condition (2) holds. (iv) Condition (3). From the bound on ηl2 above, 3456L4 η̃ 2 K 3 ρ L2 η̃ 2 3456L4 η̃ 2 K 3 ηl2 ≤ · =4 2 . 3 3 2 3 ρ ρ 864L K ρ As shown for Condition (4), this quantity is bounded by 1. Thus Condition (3) holds. A Step 3: Harmonic choice of η̃ for the first two terms. We bound the η̃–dependent part: η̃T + B η̃. Using 1/η̃ = 1/η̃0 + 1/η̃1 , A A A = + . η̃T T η̃0 T η̃1 (a) First term. Since 1/η̃0 = max{4L, 2L/ρ} ≤ 4L/ρ, we have A 4L 32L(∆F + C0 ) A ≤ · = . T η̃0 T ρ ρ2 T p (b) Second termpand linear term. Using (128), we have A/(T η̃ ) = AB/T . Also, since η̃ ≤ η̃1 , we have 1 p B η̃ ≤ B η̃1 = B A/(BT ) = AB/T . Thus r AB A . + B η̃ ≤ 2 T η̃1 T Using the definitions of A and B: r

AB = T

√ 32Lσ 2 (∆F + C0 ) 32 σ p = √ L(∆F + C0 ). 2 ρ N KT ρ N KT

√ A 32L(∆F + C0 ) 8 2σ p √ + B η̃ ≤ + L(∆F + C0 ). η̃T ρ2 T ρ N KT

(129)

Step 4: Bounding and absorbing the third term Dηl2 . By (120), ηl2 ≤

∆F + C 0 . C∗ LKσ 2 T ρ

Substituting into Dηl2 with D = 8000L2 Kσ 2 /ρ: Dηl2 ≤

8000L2 Kσ 2 ∆F + C 0 8000 L(∆F + C0 ) · = . 2 ρ C∗ LKσ T ρ C∗ ρ2 T

With C∗ = 100, this is 80L(∆F + C0 )/(ρ2 T ). Step 5: Final bound. Adding the terms: √ L(∆F + C0 ) 8 2σ p (32 + 80) + √ L(∆F + C0 ), ρ2 T ρ N KT which yields the claimed result. This corollary shows that the proposed subspace SCAFFOLD method retains the same two-term asymptotic structure as standard full-rank √ SCAFFOLD: a deterministic optimization term of order O(1/T ) and a stochastic term of order O(1/ N KT ). Hence, under the present weak assumptions—smoothness, bounded variance, and a shared random subspace—the method achieves the usual linear speedup in the total number of stochastic gradients N KT . In particular, the final rate contains no extra additive term measuring data heterogeneity; the client-drift effect is neutralized by the control variates and therefore does not worsen the asymptotic rate. 24

The only asymptotic gap relative to standard SCAFFOLD is the subspace factor: the full-rank result is recovered up to the multiplicative penalties 1/ρ in the dominant term, which quantify the price of restricting updates to an r-dimensional random subspace. Therefore, from the viewpoint of stationarity guarantees, the proposed method preserves the heterogeneity-robust convergence behavior of SCAFFOLD. This rate shows superiority over FedSub, in which the analyses rely on substantially stronger assumptions, including bounded gradients and correlation or similarity conditions between consecutive subspaces. Moreover, FedSub does not exhibit the same linear speedup; by contrast, the bound in (124) keeps the standard √ 1/ N KT term, showing that subspace compression here does not destroy the parallelization benefit of SCAFFOLD.

Experiments

This section evaluates the proposed subspace federated optimization framework on two complementary benchmarks: a controlled matrix-regression toy problem and a deep-learning benchmark on CIFAR-100 with ResNet-110. The experiments are designed to assess heterogeneity robustness, the role of the subspace dimension, and the stability of the dual mechanism. Throughout, the discussion is connected to the algorithmic design in Section 3 and the convergence intuition in Section 4. Our focus is on whether SSF can reduce communication from dimension d to dimension r while preserving as much of SCAFFOLD’s drift-correction benefit as possible.

5.1

Experimental Setup

We consider the federated matrix regression problem introduced in the method section. Each client i ∈ {1, . . . , N } holds a local dataset (Ai , Bi ) and minimizes fi (X) =

1 λ ∥Ai X − Bi ∥2F + ∥X∥2F , 2ni 2

(130)

where Ai ∈ Rni ×d , Bi ∈ Rni ×m , and X ∈ Rd×m . The global objective is N

1 X fi (X), F (X) = N i=1

(131)

and optimization quality is measured by the relative error RelErr(X) =

∥X − X ∗ ∥F , ∥X ∗ ∥F

(132)

where X ∗ denotes the closed-form minimizer of the global problem. Toy-problem configuration. Unless otherwise stated, the comprehensive toy experiments use N = 20 clients, feature dimension d = 100, output dimension m = 10, and ni = 50 samples per client. The regularization parameter is fixed to λ = 0.1 and the observation noise standard deviation is set to 0.01. Data heterogeneity is induced through client-specific feature mean shifts, µi ∼ N (0, het level2 Id ),

(133)

het level ∈ {0.1, 0.5, 2.0},

(134)

with heterogeneity levels corresponding to low, medium, and high heterogeneity, respectively. Training protocol. All four federated optimizers share the same training skeleton: S = 10 out of N = 20 clients participate in each round, every selected client performs K = 5 local steps, the minibatch size is 20, and the global step size is fixed at ηg = 1.0. The long toy experiments run for T = 25,000 communication rounds. For SSF and FedSub, the subspace ratio is chosen from r/d ∈ {0.01, 0.05, 0.10, 0.20, 0.50}, which corresponds to r ∈ {1, 5, 10, 20, 50} when d = 100. 25

(135)

Table 1: Main hyperparameters used in the federated matrix-regression experiments. Quantity

Value

Number of clients N Feature dimension d Output dimension m Samples per client ni Regularization λ Noise standard deviation Local steps K Clients per round S Batch size Global step size ηg Local step size ηℓ Heterogeneity levels Subspace ratios r/d Projector refresh (FedSub)

20 100 10 50 0.1 0.01 5 10 20 1.0 10−2 , 10−2 , 10−3 by heterogeneity 0.1, 0.5, 2.0 0.01, 0.05, 0.10, 0.20, 0.50 every 5 rounds

Learning-rate search. The local learning rate is selected by an automatic search with Full-SCAFFOLD over the candidate grid {10−4 , 10−3 , 10−2 , 10−1 } for 500 rounds (with divergence detection), selecting the rate that yields the lowest final relative error. The selected values are ηℓ = 10−2 for het = 0.1,

ηℓ = 10−2 for het = 0.5,

ηℓ = 10−3 for het = 2.0.

(136)

The chosen learning rate for each heterogeneity level is then reused for all algorithms and subspace dimensions at that level. Algorithms.

We compare four methods:

• Full-SCAFFOLD: the full-dimensional SCAFFOLD baseline; • Full-FedAvg: the classical full-dimensional FedAvg baseline; • SSF: Subspace-SCAFFOLD, which performs primal updates in a subspace while preserving fulldimensional dual information as described in Section 3; • FedSub: subspace FedAvg with subspace-only dual variables. We are particularly interested in whether SSF can remain close to Full-SCAFFOLD while reducing communication and active-device memory from order d to order r. Deep-learning benchmark. To complement the toy study, we also consider a CIFAR-100 benchmark with ResNet-110 in a heterogeneous distributed setting with four workers/clients and 100 training epochs.

5.2

Toy Experiments: Matrix Regression

The toy benchmark is especially useful because the global optimum is available in closed form, so every curve can be interpreted directly in terms of optimization error rather than proxy metrics. We organize the analysis around three questions: robustness to data heterogeneity, dependence on subspace dimension, and the stability of FedSub.

(a) Low heterogeneity (het = 0.1, r = (b) Medium heterogeneity (het = 0.5, (c) High heterogeneity (het = 2.0, r = 20). r = 20). 20).

Figure 1: Relative error versus communication rounds for r = 20 under low, medium, and high heterogeneity. The y-axis shows the relative error ∥Xt − X ∗ ∥F /∥X ∗ ∥F on a logarithmic scale. Table 2: Final relative error at T = 25,000 communication rounds for fixed subspace dimension r = 20 and varying heterogeneity.

Heterogeneity

Name

Full-SCAFFOLD

Full-FedAvg

SSF

FedSub

0.1 0.5 2.0

low medium high

20 20 20

6.9726e-03 6.4701e-03 2.0831e-03

9.1265e-03 1.1197e-02 3.8143e-03

7.5535e-03 8.2431e-03 3.4495e-03

1.1295e-02 1.0556e-02 6.0512e-03

5.2.1

Robustness to data heterogeneity at fixed subspace dimension

We first fix the subspace dimension to r = 20 (r/d = 0.20) and vary the heterogeneity level het level ∈ {0.1, 0.5, 2.0}. Figure 1 shows the corresponding convergence curves, and Table 2 reports the final relative errors at the end of training. For low heterogeneity (het = 0.1; Figure 1a), all four algorithms remain numerically stable and converge to the 10−3 –10−2 range. Full-SCAFFOLD reaches 6.9726e-03, SSF reaches 7.5535e-03, Full-FedAvg reaches 9.1265e-03, and FedSub is highest at 1.1295e-02. For medium heterogeneity (het = 0.5; Figure 1b), Full-SCAFFOLD is best (6.4701e-03), SSF is second (8.2431e-03), and FedSub/Full-FedAvg finish at 1.0556e-02 and 1.1197e-02, respectively. For high heterogeneity (het = 2.0; Figure 1c), all methods still converge stably, but the full-dimensional advantage is clearer: Full-SCAFFOLD reaches 2.0831e-03, Full-FedAvg reaches 3.8143e-03, SSF reaches 3.4495e-03, and FedSub reaches 6.0512e-03. Overall, for r/d = 0.2 and the heterogeneity split (het = 0.1/0.5/2.0 for low/medium/high), SSF remains competitive with full-dimensional training and consistently outperforms FedSub. 5.2.2

Effect of the subspace dimension at high heterogeneity

We next fix heterogeneity at het = 2.0 (high under the current split) and vary r ∈ {1, 5, 10, 20, 50}. Figure 2 presents representative convergence curves, and Table 3 summarizes the final errors. As expected, the full-dimensional methods are nearly insensitive to the choice of r: Full-SCAFFOLD and Full-FedAvg achieve essentially constant final errors in the ranges [2.03, 2.12]e-03 and [3.65, 3.89]e-03, respectively. For SSF, the subspace dimension plays a crucial role. At r = 1 (1% of the ambient dimension), SSF plateaus at 2.82e-01, more than two orders of magnitude worse than Full-SCAFFOLD. Increasing the dimension to r = 5 already yields a dramatic improvement: the error drops to 7.81e-03, nearly 36 times smaller. For r ∈ {10, 20, 50}, SSF continues to improve and its final error lies within a small factor of the full-dimensional optimum, e.g., 3.03e-03 at r = 50 versus 2.08e-03 for Full-SCAFFOLD.

(a) r = 10 (r/d = 0.10).

(b) r = 20 (r/d = 0.20).

Figure 2: Effect of the subspace dimension on convergence at high heterogeneity (het = 2.0). Full-SCAFFOLD and Full-FedAvg are insensitive to r, whereas the behaviour of SSF and FedSub depends strongly on r. Table 3: Final relative error at T = 25,000 communication rounds for high heterogeneity (het = 2.0) and varying subspace dimension. The row r = 50 highlights the numerical divergence of FedSub. r/d

Full-SCAFFOLD

Full-FedAvg

SSF

FedSub

0.01 0.05 0.10 0.20 0.50

1 5 10 20 50

2.12e-03 2.03e-03 2.06e-03 2.08e-03 2.08e-03

3.65e-03 3.71e-03 3.89e-03 3.81e-03 3.87e-03

2.82e-01 7.81e-03 3.74e-03 3.45e-03 3.03e-03

2.81e-01 6.76e-03 4.32e-03 6.05e-03 NaN (diverged)

FedSub exhibits a non-monotone and less stable dependence on r. At r = 1 it behaves similarly to SSF, with error 2.81e-01. At r = 5 and r = 10, FedSub improves substantially (6.76e-03 and 4.32e-03, respectively) and becomes comparable to SSF. However, at r = 20 its performance degrades again (6.05e-03), and at r = 50 the run becomes numerically unstable: the CSV reports a final error of NaN, indicating divergence. This pattern directly supports the mechanism described in Section 3. Because SSF preserves residual dual information in the orthogonal complement of the active projector, enlarging r improves the quality of the primal subspace updates without discarding the full-space drift information. FedSub, in contrast, keeps only subspace-only dual states, so increasing the subspace dimension does not eliminate the structural information loss discussed in Section 3; in fact, at large r the method can become unstable rather than reliably approaching the full-dimensional behavior. 5.2.3

FedSub stability at large subspace dimension

To isolate the instability phenomenon, we examine the r = 50 setting across all heterogeneity levels. Table 4 shows that the final FedSub error is recorded as NaN for all three heterogeneity levels, while Figure 3 illustrates a representative medium-heterogeneity case. Across all three heterogeneity levels, Full-SCAFFOLD and Full-FedAvg remain stable. SSF is also consistently stable, with final errors 7.69e-03 (low, het = 0.1), 8.70e-03 (medium, het = 0.5), and 3.03e-03 (high, het = 2.0), remaining within roughly one order of magnitude of the full-dimensional baselines. In stark contrast, FedSub fails in all three cases at r = 50: the final error is reported as NaN regardless of the heterogeneity level, indicating numerical instability or divergence during training. The instability is also apparent in the convergence plots: the FedSub curve displays a large transient spike to extremely high error and does not settle into a stable plateau comparable to the other methods. These empirical findings align with the theoretical discussion in Section 4. In FedSub, the control variates (dual variables) are repeatedly projected into the low-rank subspace, discarding information in the orthogonal complement. Under heterogeneous data and large r, this projection prevents the effective control variate from converging to a stationary point of the full gradient field, leading to drifting or exploding dual variables and numerical instability. In contrast, SSF maintains full-dimensional dual variables while restricting only

Table 4: Final relative error at T = 25,000 communication rounds for subspace dimension r = 50 (r/d = 0.5) and varying heterogeneity. FedSub diverges (numerically) in all three cases.

Heterogeneity

Name

Full-SCAFFOLD

Full-FedAvg

SSF

FedSub

0.1 0.5 2.0

low medium high

50 50 50

7.27e-03 7.00e-03 2.08e-03

8.25e-03 1.04e-02 3.87e-03

7.69e-03 8.70e-03 3.03e-03

NaN (diverged) NaN (diverged) NaN (diverged)

Figure 3: Convergence behaviour at medium heterogeneity (het = 0.5) and large subspace dimension r = 50. While Full-SCAFFOLD, Full-FedAvg, and SSF converge smoothly, FedSub exhibits pronounced instability with large error excursions, and the final recorded error is NaN. the primal updates to the subspace, preserving the residual information in the orthogonal complement and yielding much more robust behaviour.

5.3

Deep-Learning Benchmark: CIFAR-100 with ResNet-110

The second experimental component uses the CIFAR-100 benchmark with the final submission logs. Figure 4 plots test accuracy over training, and Table 5 summarizes the final and best accuracies over 100 epochs. The deep-learning ranking is unambiguous: Full-SCAFFOLD > SSF > FedAvg > FedSub.

(137)

The margin between Full-SCAFFOLD and SSF is approximately 7.5 percentage points at the final epoch, while SSF substantially outperforms FedAvg and FedSub. Notably, in this submission run FedSub underperforms FedAvg, indicating that the subspace FedAvg-style control mechanism is sensitive to the concrete training protocol and may require additional tuning for deep models. We also note that the current submission protocol is not strictly equal-compute across all methods: Full-SCAFFOLD follows a different epoch-internal schedule than the other three methods in the logs. Therefore, these results should be interpreted as submission-level empirical performance rather than a fully compute-matched ablation. A strict equal-budget comparison is left for follow-up experiments.

Figure 4: Test accuracy versus training epoch on CIFAR-100 with ResNet-110. Full-SCAFFOLD is the strongest method, SSF is consistently second best, and both FedAvg and FedSub trail by a clear margin. This figure reports test performance only. Table 5: CIFAR-100/ResNet-110 results over 100 epochs. “Final” denotes epoch 99; “Best” denotes the best test accuracy during training. Method Full-SCAFFOLD SSF FedAvg FedSub

5.4

Final Test Acc. (%)

Best Test Acc. (%)

Best Epoch

52.92 45.43 34.35 23.17

53.11 45.43 38.51 30.20

53 99 46 41

Discussion

Combining the toy and deep-learning experiments yields four conclusions. First, method robustness under controlled heterogeneity: in the toy benchmark with the heterogeneity split (het = 0.1/0.5/2.0 for low/medium/high), SSF is competitive with full-dimensional baselines at moderate subspace dimension (r/d = 0.2), while FedSub remains less robust. Second, subspace-dimension sensitivity: SSF improves consistently as r increases and approaches fulldimensional performance for r ≥ 10, whereas FedSub exhibits non-monotone behavior and instability at large r. Third, FedSub stability limitations: FedSub diverges in multiple large-subspace toy settings and also underperforms FedAvg in the submitted CIFAR-100 run, indicating limited robustness of repeatedly projected control variates under heterogeneous optimization dynamics. Fourth, practical deep-learning outcome: in the final submission setting on CIFAR-100/ResNet-110, Full-SCAFFOLD is the strongest method and SSF is the second-best method with a clear gain over FedAvg and FedSub. This confirms that SSF is a practically effective compromise between full-dimensional control and subspace efficiency, while still leaving headroom to Full-SCAFFOLD.

Conclusion

This paper addresses a central limitation of large-model federated learning: existing methods typically achieve either robustness to statistical heterogeneity or efficiency in communication, memory, and computation, but not both simultaneously. Motivated by this gap, we develop SSF, which brings heterogeneity-correction into a low-dimensional optimization subspace and preserves the accumulated correction information when the subspace changes. This subspace-native design directly targets the large-model FL regime highlighted in the introduction and positions SSF as a bridge between full-dimensional heterogeneity-correction methods and efficient subspace training. The main findings are twofold. On the theoretical side, under standard smoothness and bounded-variance assumptions, SSF preserves the bounded-variance structure of projected stochastic gradients and achieves a √ e non-asymptotic convergence rate of order O(1/T + 1/ N KT ), matching the linear speedup behavior expected from strong federated heterogeneity-correction methods while avoiding full-dimensional auxiliary states. On the empirical side, the experiments indicate that SSF delivers favorable accuracy–efficiency trade-offs in heterogeneous federated settings, reducing memory usage, communication volume, and local computation while maintaining robustness to client drift. Taken together, these results suggest that native low-dimensional heterogeneity correction is a promising direction for scalable and practical federated optimization.

Acknowledgments We gratefully acknowledge ReasFlow [19], a reasoning-centric scientific discovery assistant, for its substantial contributions to the preparation of this paper. A significant portion of the work, including the literature review, mathematical proofs, numerical experiments, and the initial manuscript draft, was generated automatically with the assistance of ReasFlow. The authors’ contributions lay primarily in identifying the research problem, proposing the high-level algorithmic design, articulating the key ideas underlying the mathematical proofs, specifying the methodology and requirements for the numerical experiments, and polishing the manuscript to meet the standards required for submission. In particular, the authors devoted considerable effort to verifying the correctness of the mathematical proofs and refining the resulting arguments.

References [1] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communicationefficient SGD via gradient quantization and encoding. arXiv preprint arXiv:1610.02132v4, 2016. [2] Yajie Bao, Michael Crawshaw, Shan Luo, and Mingrui Liu. Fast composite optimization and statistical recovery in federated learning. In Proceedings of the International Conference on Machine Learning, 2022. [3] Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-SGD: Distributed SGD with quantization, sparsification, and local computations. IEEE Journal on Selected Areas in Information Theory, 2019. [4] Debora Caldarola, Pietro Cagnasso, Barbara Caputo, and Marco Ciccone. Beyond local sharpness: Communication-efficient global sharpness-aware minimization for federated learning. 2024. arXiv:2412.03752v2. [5] Chuyan Chen, Yutong He, Pengrui Li, Weichen Jia, and Kun Yuan. Greedy low-rank gradient compression for distributed learning with convergence guarantees. arXiv preprint arXiv:2507.08784v4, 2025. [6] Ilyas Fatkhullin, Alexander Tyurin, and Peter Richtárik. Momentum provably improves error feedback! In Advances in Neural Information Processing Systems, 2023. [7] Xinmeng Huang, Ping Li, and Xiaoyun Li. Stochastic controlled averaging for federated learning with communication compression. In International Conference on Learning Representations, 2023.

[8] Like Jian and Dong Liu. Widening the network mitigates the impact of data heterogeneity on FedAvg. In Proceedings of the 41st International Conference on Machine Learning, 2025. [9] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. arXiv preprint arXiv:1910.06378, 2020. [10] Anastasia Koloskova, Tao Lin, and Sebastian U. Stich. An improved analysis of gradient tracking for decentralized machine learning. In Advances in Neural Information Processing Systems, volume 34, pages 11452–11465, 2021. [11] Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492v2, 2016. [12] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on non-IID data. In International Conference on Learning Representations, 2020. [13] Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. In International Conference on Learning Representations, 2017. [14] Junkang Liu, Fanhua Shang, Junchao Zhou, Hongying Liu, Yuanyuan Liu, and Jin Liu. Fedmuon: Accelerating federated learning with matrix orthogonalization. arXiv preprint arXiv:2510.27403v1, 2025. [15] Yue Liu, Tao Lin, Anastasia Koloskova, and Sebastian U. Stich. Decentralized gradient tracking with local steps. Optimization Methods and Software, 2023. [16] H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Ag”uera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, pages 1273–1282, 2017. [17] Zhaonan Qu, Kaixiang Lin, Zhaojian Li, Jiayu Zhou, and Zhengyuan Zhou. A unified linear speedup analysis of federated averaging and nesterov fedavg. Journal of Artificial Intelligence Research, 76:1345– 1386, 2023. [18] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In International Conference for High Performance Computing, Networking, Storage and Analysis, 2019. [19] ReasFlow Team. Reasflow: Assisting reasoning-centric scientific discovery in applied mathematics via a knowledge-based multi-agent system, 2026. [20] Thomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from low-dimensional gradient statistics. In International Conference on Learning Representations, 2024. [21] Sebastian U. Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019. [22] Sebastian U. Stich and Sai Praneeth Karimireddy. The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication. arXiv preprint arXiv:1909.05350v2, 2019. [23] Hanlin Tang, Xiangru Lian, Chen Yu, Tong Zhang, and Ji Liu. Doublesqueeze: Parallel stochastic gradient descent with double-pass error-compensated compression. In Proceedings of the International Conference on Machine Learning, 2019. [24] Sahil Tyagi and Martin Swany. Gravac: Adaptive compression for communication-efficient distributed dl training. In IEEE International Conference on Cloud Computing, 2023.

[25] Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. Gradient sparsification for communication-efficient distributed optimization. In Advances in Neural Information Processing Systems, 2017. [26] Ming Xiang, Stratis Ioannidis, Edmund Yeh, Carlee Joe-Wong, and Lili Su. Towards bias correction of FedAvg over nonuniform and time-varying communications. In Proceedings of the IEEE Conference on Decision and Control, pages 5815–5822, 2023. [27] Haibo Yang, Minghong Fang, and Jia Liu. Achieving linear speedup with partial worker participation in non-IID federated learning. In International Conference on Learning Representations, 2021. [28] Fuxun Yu, Weishan Zhang, Zhuwei Qin, Zirui Xu, Di Wang, Chenchen Liu, Zhi Tian, and Xiang Chen. Heterogeneous federated learning. arXiv preprint arXiv:2008.06767, 2020. [29] Honglin Yuan and Tengyu Ma. Federated accelerated stochastic gradient descent. In Advances in Neural Information Processing Systems, 2020. [30] Jiaojiao Zhang, Jiang Hu, and Mikael Johansson. Non-convex composite federated learning with heterogeneous data. arXiv preprint arXiv:2502.03958v1, 2025. [31] Jiaojiao Zhang, Yuqi Xu, and Kun Yuan. An efficient subspace algorithm for federated learning on heterogeneous data. arXiv preprint arXiv:2509.05213v1, 2025. [32] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. OPT: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068v4, 2022. [33] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. In Proceedings of the International Conference on Machine Learning, 2024. [34] Yuan Zhou, Jiachen Zhong, Xinli Shi, Guanghui Wen, and Xinghuo Yu. FedCanon: Non-convex composite federated learning with efficient proximal operation on heterogeneous data. IEEE Transactions on Signal Processing, 2025.

Record · ID 141497 · SHA-256 3d5558addc91a81c
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.