cs.LGOct 2, 2026

Distributed Learning with Selective State Space Models: Architecture-Aware Convergence Analysis

Authors: Adam Piaseczny, Md Kamran Chowdhury Shisher, Shiqiang Wang, Christopher G. Brinton

Organizations: Purdue University · University of Colorado Colorado Springs · University of Exeter

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Modern state space models (SSMs), such as Mamba2, provide a compelling alternative to transformers by combining linear-time sequence modeling with recurrent state-space dynamics. However, the behavior of SSMs in distributed learning settings remains poorly understood. In particular, the existing standard federated learning methods are largely architecture-agnostic, and do not account for the stability, selectivity, and state-space parameterization that characterize modern selective SSMs. To address this, we derive architecture-aware gradient and smoothness bounds for single- and multi-layer selective SSMs, and convergence bounds for FedAvg and FedProx, characterizing how recurrent stability, input-dependent discretization, and state projection norms affect federated optimization. We then numerically validate the single-layer bounds on sequences generated by a teacher SSM, using a learner that follows the analyzed recurrence. We use this analysis to formulate expectations about the effects of local training and client heterogeneity, and examine these expectations by comparing nine federated learning algorithms on Mamba2 language modeling across six text domains. These experiments illustrate how SSM-specific bounds can provide a basis for interpreting the behavior of practical federated learning algorithms.

Figures & tables

Figure 1: (a) A selective SSM layer generates input-dependent discretization parameters from xtx_{t} and uses them to update the hidden state hth_{t} and output yty_{t} . (b) Federated SSM training proceeds over communication rounds: the server broadcasts WrW^{r} , clients perform EE local updates on local datasets Dk\mathcal{D}_{k} , and the resulting models are aggregated into Wr+1W^{r+1} .
StatementInequality checkedEvaluationsViolationsMax. utilization
Thm. 4.4∥∇L∥F≤GSSM\|\nabla\mathcal{L}\|_{F}\leq\mathcal{G}_{\mathrm{SSM}} , per sequence1,327,1040 / 01.0×10−31.0\times 10^{-3} / 8.3×10−38.3\times 10^{-3}
Lemma 4.7∥∇2L∥2^≤LSSM,0\widehat{\|\nabla^{2}\mathcal{L}\|_{2}}\leq L_{\mathrm{SSM},0}41,47207.0×10−57.0\times 10^{-5}
Lemma 4.7finite-difference ratio ≤LSSM,0.001\leq L_{\mathrm{SSM},0.001}41,47208.2×10−58.2\times 10^{-5}
Lemma 4.7∥∇L(Wa)−∇L(Wb)∥F≤LSSM,R∥Wa−Wb∥F\|\nabla\mathcal{L}(W_{a})-\nabla\mathcal{L}(W_{b})\|_{F}\leq L_{\mathrm{SSM},R}\|W_{a}-W_{b}\|_{F}86401.0×10−51.0\times 10^{-5}
Lemma B.1∣GSSM(Wa)−GSSM(Wb)∣≤LGSSM,R∥Wa−Wb∥F|\mathcal{G}_{\mathrm{SSM}}(W_{a})-\mathcal{G}_{\mathrm{SSM}}(W_{b})|\leq L_{\mathcal{G}_{\mathrm{SSM}},R}\|W_{a}-W_{b}\|_{F}86400.420.42
Table 1: Gradient and smoothness bounds on the synthetic bank: 1,411,776 evaluations, no violations detected. Maximum utilization is the largest measured-to-derived ratio, with a/ba/b denoting configured/tightened constants. Evaluations count test sequences at each checkpoint (first row), sequence-checkpoint pairs (curvature rows), or consecutive-checkpoint pairs (trajectory rows).
Figure 2: Single-layer bounds on the SSM teacher-generated datasets. ( 2(a) ) Across 1,296 learner and teacher parameter points, every measured per-sequence gradient norm lies below the configured GSSM\mathcal{G}_{\mathrm{SSM}} . ( 2(b) ) At training end, GSSM\mathcal{G}_{\mathrm{SSM}} and LSSM,RL_{\mathrm{SSM},R} scale as (1−ρ)−1.8(1-\rho)^{-1.8} and (1−ρ)−3.6(1-\rho)^{-3.6} , versus (1−ρ)−0.5(1-\rho)^{-0.5} and (1−ρ)−1.0(1-\rho)^{-1.0} for measured gradient and estimated Hessian norms. ( 2(c) ) As ρ→1\rho\to 1 , gradient energy shifts from WB,WCW_{B},W_{C} toward the transition block AA . ( 2(d) ) One-factor fits estimate the power-law exponent α\alpha for the derived (filled) and measured (open) gradient norms as vocabulary size DD (blue) or state size NN (orange) varies, with bootstrap 95% intervals.
  ExpectedΔ\Delta PPL (%)
MethodTheory-linked mechanismβ=0\beta=0β=1\beta=1β=0\beta=0β=1\beta=1
FedProx ( Li et al., 2020 )radius, Rμ≤RR_{\mu}\leq R−-−-−17.0-17.0−34.4-34.4
SCAFFOLD ( Karimireddy et al., 2020 )corrected disagreement++−-+4.0+4.0−19.4-19.4
FedDyn ( Acar et al., 2021 )heterogeneity; radius (proximal α\alpha )−-−-−2.2-2.2−35.4-35.4
FedAlign ( Keçeci et al., 2025 )state-coordinate alignment====−0.3-0.3−1.2-1.2
FedSAM ( Qu et al., 2022 )expanded evaluation neighborhood====−0.4-0.4−1.7-1.7
Table 2: Theorem-motivated expectations versus observed behavior on Mamba2. −- / == / ++ denote expected lower/similar/higher validation perplexity relative to FedAvg, motivated by the listed mechanisms. Δ\Delta PPL is the mean percentage change across the 2424 configurations at each β\beta .
Appendix figures & tables38 assets

Supplementary material from the paper’s appendix.

Appendix

SettingValue
Input/vocabulary dimension DD{4,8,16}\{4,8,16\}
State dimension NN{4,8,16}\{4,8,16\}
Stability limit ρ\rho{0.3,0.5,0.7,0.9,0.95,0.99,0.995,0.999}\{0.3,0.5,0.7,0.9,0.95,0.99,0.995,0.999\}
Configurations (datasets)7272
Sequence length SS6464 stored tokens ( 1616 burn-in tokens discarded, 8080 generated)
Training / test sequences4,0964{,}096 / 1,0241{,}024
Appendix
Table 3: Generation settings of the synthetic teacher datasets. Grid entries are searched per configuration in the fixed order given in the text.
CriterionThreshold
Global token balancelargest probability in the token marginal ≤1/D+0.10\leq 1/D+0.10
Marginal entropyShannon entropy of the token marginal ≥0.90ln⁡D\geq 0.90\ln D (the maximum ln⁡D\ln D is attained by the uniform distribution)
Per-sequence majority (test split)for each test sequence, the fraction of its SS positions occupied by its most frequent token; the 9595 th percentile of this fraction over the test sequences ≤1/D+0.20\leq 1/D+0.20
Class coverageeach of the DD tokens occurs at least once in the training split and at least once in the test split
Constant sequencesfraction of sequences whose SS tokens are all identical ≤0.01\leq 0.01
Sequence diversitynumber of distinct sequences divided by the number of sequences ≥0.95\geq 0.95
Appendix
Table 4: Acceptance criteria applied to every synthetic candidate dataset. Statistics are computed on the union of the training and test splits unless noted. The token marginal is the empirical distribution of the DD tokens over all stored positions, ignoring position and context. NLL denotes the teacher-forced negative log-likelihood at the sampling temperature τ\tau with h0=0h_{0}=\mathbf{0} per stored sequence.
Settingβ\betaTrain capVal capTrain limiter dataValid limiter dataServer val
Setting 10.07,4271,857stackexchangestackexchange7,424
Setting 10.77,4271,857stackexchangestackexchange7,431
Setting 11.07,4271,857stackexchangestackexchange7,428
Setting 20.07,3671,842c4c47,360
Setting 20.77,3671,842c4c47,364
Setting 21.07,3671,842c4c47,368
Appendix
Table 5: Effective client train/validation budgets and server validation sizes across heterogeneity settings.
DatasetTask# ClassesTrain capEval capmax_seq_len
ag_newsNews topic440,0005,000256
dbpedia_14Wikipedia topic1440,0005,000256
trecQuestion type6≈5,500\approx 5{,}500 (full)5,000256
imdbSentiment240,0005,000256
Appendix
Table 6: Task datasets and sequence limits for causal-LM evaluation.
DatasetPrompt template and labels
ag_newsArticle: {text}\nCategory:
 Labels: World , Sports , Business , Tech
dbpedia_14<title-aware template>...:Category:
 Labels: 14-way Wikipedia topic labels
trecQuestion: {text}\nType:
 Labels: Description , Entity , Abbreviation , Person , Location , Number
Appendix
Table 7: Prompt templates and candidate labels used for causal-LM scoring.
 N=4N=4N=8N=8N=16N=16
D=4D=452528888160160
D=8D=8132132200200336336
D=16D=16388388520520784784
Appendix
Table 8: Learner parameter counts D2+2ND+ND^{2}+2ND+N for the synthetic configurations.
DepthTotal paramsSSM (mixer) paramsEmbedding paramsPer-layer mixer
414,733,6641,860,70412,871,680465,176
1620,318,8487,442,81612,871,680465,176
Appendix
Table 9: Parameter breakdown by model depth.
BackboneLayoutSSM fractionTotal params
mamba24×4\times Mamba21.001.0014,733,66414{,}733{,}664
hybrid75attention at {2}\{2\}0.750.7514,531,65614{,}531{,}656
hybrid50attention at {1,3}\{1,3\}0.500.5014,329,64814{,}329{,}648
hybrid50altattention at {0,2}\{0,2\}0.500.5014,329,64814{,}329{,}648
hybrid25attention at {0,2,3}\{0,2,3\}0.250.2514,127,64014{,}127{,}640
attn04×4\times attention0.000.0013,925,63213{,}925{,}632
Appendix
Table 10: Backbones of the depth- 44 ablation. Layer indices run from 00 to 33 , every layer not listed as attention is a Mamba2 mixer.
CheckpointTotal params
state-spaces/mamba2-130m130M
state-spaces/mamba2-370m370M
state-spaces/mamba2-780m780M
Appendix
Table 11: Public Mamba2 checkpoints used for model merging.
ComponentValue
Lossmean softmax cross-entropy over the S−1S-1 next-token positions
OptimizerSGD, no momentum, weight decay, or gradient clipping
Learning rate η\eta0.050.05
Batch size6464 sequences
Centralized updates2,0002{,}000 (epochs of 6464 updates, reshuffled every epoch)
Checkpoints and evaluationevery 100100 updates; bound and smoothness
Appendix
Table 12: Shared settings of the synthetic centralized and federated runs.
ComponentValue
AlgorithmsFedAvg; FedProx
Base settingK=4K=4 clients, E=16E=16 local updates, IID partition
Communication rounds TT6060
Client-count variationK∈{2,8,32}K\in\{2,8,32\}
Local-step variationE∈{1,4,64}E\in\{1,4,64\}
Proximal variation (FedProx)μ∈{0.01,1}\mu\in\{0.01,1\}
Appendix
Table 13: Synthetic federated settings. The base setting is run on all 7272 configurations, each variation changes one factor of the base setting and is run on the 1616 core configurations.
ComponentValue
Number of clients, KK44
Participation rate1.01.0 ; all clients participate each round
Communication rounds, TT5050
Local steps per round, EE{64,128,256}\{64,128,256\}
Per-step batch size88
Sequence length, SS512512
Appendix
Table 14: Shared training hyperparameters used across all federated algorithms.
AlgorithmHyperparameters and values
FedAvgnone
FedProxμ∈{5×10−5, 10−4, 5×10−4, 10−3, 5×10−2}\mu\in\{5{\times}10^{-5},\,10^{-4},\,5{\times}10^{-4},\,10^{-3},\,5{\times}10^{-2}\}
SCAFFOLDfc∈{4,8,16}f_{c}\in\{4,8,16\} ; ηc=0.01\eta_{c}=0.01 ; λc=0.01\lambda_{c}=0.01 ; τc=0.1\tau_{c}=0.1
FedAdamαs∈{10−3, 5×10−3, 10−2}\alpha_{s}\in\{10^{-3},\,5{\times}10^{-3},\,10^{-2}\} ; β1s=0.9\beta_{1}^{s}=0.9 ; β2s=0.99\beta_{2}^{s}=0.99 ; τs=10−3\tau_{s}=10^{-3}
FedAlignalign_mode == mamba2_ref ; align_mamba2_heads ∈{0,1}\in\{0,1\}
FedDynα∈{10−3, 3×10−3, 10−2, 3×10−2, 10−1}\alpha\in\{10^{-3},\,3{\times}10^{-3},\,10^{-2},\,3{\times}10^{-2},\,10^{-1}\}
Appendix
Table 15: Hyperparameter grids of the federated algorithms, run at every configuration and seed. Values shared by all members of a grid are fixed.
ComponentValue
OptimizerAdamW
Learning rate2×10−42\times 10^{-4}
Weight decay0.0
Gradient clippingℓ2\ell_{2} norm, threshold 1.0
Train batch size4 for mamba2-130m, mamba2-370m , 8 for mamba2-780m
Eval batch size4
Appendix
Table 16: Per-task fine-tuning recipe (identical across tasks and seeds).
ρ\rhocellspointssequencesviol. (conf./tight.)median util. (conf.)median util. (tight.)max util. (tight.)median ρ^\hat{\rho}
0.3 ∗11818,4320 / 07.6×10−57.6\times 10^{-5}7.2×10−47.2\times 10^{-4}8.3×10−38.3\times 10^{-3}0.0602
0.38144147,4560 / 02.0×10−62.0\times 10^{-6}1.2×10−41.2\times 10^{-4}4.1×10−34.1\times 10^{-3}0.0345
0.59162165,8880 / 01.6×10−61.6\times 10^{-6}3.7×10−53.7\times 10^{-5}5.1×10−35.1\times 10^{-3}0.146
0.79162165,8880 / 07.7×10−77.7\times 10^{-7}1.6×10−51.6\times 10^{-5}2.1×10−32.1\times 10^{-3}0.369
0.99162165,8880 / 02.3×10−72.3\times 10^{-7}4.1×10−64.1\times 10^{-6}1.0×10−31.0\times 10^{-3}0.746
0.959162165,8880 / 09.1×10−89.1\times 10^{-8}2.0×10−62.0\times 10^{-6}5.7×10−45.7\times 10^{-4}0.868
Appendix
Table 17: Validity and tightness of the gradient bound of Theorem 4.4 by stability limit (single layer). Utilization is the measured per-sequence gradient norm divided by GSSM\mathcal{G}_{\mathrm{SSM}} (conf.) or GSSMtight\mathcal{G}^{\mathrm{tight}}_{\mathrm{SSM}} (tight.). The row marked ∗ is the one ρ=0.3\rho=0.3 cell whose realized step δ^\hat{\delta} differs from the other eight and is grouped separately by the evaluation.
checkpointpointsmedian util. (conf.)median util. (tight.)max util. (tight.)median GSSM\mathcal{G}_{\mathrm{SSM}}median gradient norm
frozen_teacher2162.7×10−82.7\times 10^{-8}4.0×10−74.0\times 10^{-7}1.8×10−41.8\times 10^{-4}9.0×1069.0\times 10^{6}0.177
learner_final2161.2×10−71.2\times 10^{-7}4.4×10−64.4\times 10^{-6}3.5×10−33.5\times 10^{-3}6.5×1056.5\times 10^{5}0.0612
learner_initialization2163.4×10−73.4\times 10^{-7}1.9×10−51.9\times 10^{-5}8.3×10−38.3\times 10^{-3}1.1×1051.1\times 10^{5}0.0293
learner_intermediate6481.4×10−71.4\times 10^{-7}5.8×10−65.8\times 10^{-6}7.4×10−37.4\times 10^{-3}4.0×1054.0\times 10^{5}0.0438
Appendix
Table 18: Utilization along training (single layer).
Figure 3: Median bound utilization along training, one curve per ρ\rho (colour), with the frozen teacher at the last position. The bound is tightest at initialization and loosest at the frozen teacher at every ρ\rho .
Figure 4: ( 4(a) ) Largest measured gradient norm over the test split against the tightened GSSMtight\mathcal{G}^{\mathrm{tight}}_{\mathrm{SSM}} for all 1,296 parameter points (thin lines mark utilizations 10−210^{-2} to 10−610^{-6} ). ( 4(b) ) Empirical distribution of the per-sequence tightened utilization per ρ\rho (one in eight sequences shown), the vertical line is the bound.
checkpointρ\rhoWΔW_{\Delta}WBW_{B}WCW_{C}AArecurrent share of GSSM\mathcal{G}_{\mathrm{SSM}}
final (step 2000)0.30.1480.4280.4310.0000.994
 0.50.1450.4340.4330.0000.996
 0.70.1210.4320.4280.0020.997
 0.90.0440.3690.4150.0390.999
 0.950.0390.3870.4290.0790.999
 0.990.0060.1870.2040.6001.000
Appendix
Table 19: Where the measured gradient lives (median share of squared Frobenius norm per parameter block) against the share of the certified bound carried by its recurrent-Jacobian term.
ρ\rhoG\mathcal{G} in DDGSSMtight\mathcal{G}^{\mathrm{tight}}_{\mathrm{SSM}} in DD∥∇L∥\|\nabla\mathcal{L}\| in DDG\mathcal{G} in NNGSSMtight\mathcal{G}^{\mathrm{tight}}_{\mathrm{SSM}} in NN∥∇L∥\|\nabla\mathcal{L}\| in NN
0.31.78 [1.71, 2.84]-0.29 [-1.28, 0.11]-0.92 [-0.95, 0.41]1.96 [1.81, 2.10]1.81 [0.82, 2.30]0.17 [0.14, 0.22]
0.51.75 [1.48, 2.80]0.38 [-0.79, 0.77]-0.97 [-1.09, 0.34]2.05 [1.76, 2.21]1.86 [1.21, 2.44]0.19 [0.12, 0.29]
0.71.39 [1.30, 2.47]0.04 [-0.40, 0.89]-1.23 [-1.49, -0.07]1.70 [0.94, 1.80]1.41 [1.32, 1.96]0.03 [-0.37, 0.27]
0.91.86 [1.70, 2.19]0.48 [0.33, 1.64]-1.28 [-1.66, -1.09]1.32 [0.74, 2.23]1.55 [0.30, 3.18]0.10 [-0.50, 0.87]
0.951.74 [1.25, 2.01]0.31 [0.10, 1.40]-1.63 [-1.84, -1.58]1.09 [0.87, 1.14]1.37 [0.28, 1.55]0.05 [-0.07, 0.23]
0.991.85 [0.88, 2.23]1.05 [0.03, 1.96]-1.38 [-1.58, -1.27]1.08 [0.95, 1.11]1.38 [0.75, 1.66]0.06 [-0.12, 0.12]
Appendix
Table 20: Log-log exponents in the state and vocabulary sizes at the final learner checkpoint (median over one-factor fits with 95% confidence intervals).
ρ\rhorunsrecovered fractionmedian excess NLLmedian gap recovered [95% CI]median final ρ^\hat{\rho}
0.3270.930.01800.80 [0.52, 0.90]0.031
0.5270.930.02120.79 [0.49, 0.82]0.134
0.7270.930.02030.81 [0.40, 0.87]0.357
0.9270.930.01870.69 [0.64, 0.96]0.724
0.95271.000.01790.74 [0.52, 0.99]0.855
0.99271.000.00270.95 [0.42, 1.03]0.966
Appendix
Table 21: Centralized learners by stability limit: fraction of runs whose test Negative Log Likelihood (NLL) falls below the uniform baseline, median excess NLL over the teacher, and the fraction of the uniform-to-teacher gap recovered.
Figure 5: ( 5(a) ) Fraction of the uniform-to-teacher NLL gap recovered by the learners, per ρ\rho . ( 5(b) ) The teachers’ predictive gap over the uniform baseline as a percentage of the uniform NLL, per cell.
Figure 6: Federated scaling directions (medians with bootstrap 95% intervals). ( 6(a) ) End-of-round client drift against EE for FedAvg and FedProx on the core cells, with the initialization-based reference radius R=ηE GSSM(W0)R=\eta E\,\mathcal{G}_{\mathrm{SSM}}(W^{0}) . ( 6(b) ) Client dissimilarity ∑kpk∥∇Lk(Wr)−∇f(Wr)∥F2\sum_{k}p_{k}\|\nabla\mathcal{L}_{k}(W^{r})-\nabla f(W^{r})\|_{F}^{2} against KK on the core cells. ( 6(c) ) ∑kpk∥∇Lk(Wr)−∇f(Wr)∥F2\sum_{k}p_{k}\|\nabla\mathcal{L}_{k}(W^{r})-\nabla f(W^{r})\|_{F}^{2} , drift and ∥∇f∥F2\|\nabla f\|_{F}^{2} against 1−ρ1-\rho at the base setting (all 72 cells).
Figure 7: Per-round gradient (top) and client drift (bottom) on the core cells for each EE
Figure 8: ( 8(a) ) Measured client dissimilarity against the certified heterogeneity bound with the tightened constant, for IID and dominant-token splits. ( 8(b) ) Sequence-level, bigram and unigram weighted total variation for the client counts and partitions of the bank, only the sequence-level value enters the lemma.
factorrange (decades)std (decades)share of Var(log⁡10GSSM)\mathrm{Var}(\log_{10}\mathcal{G}_{\mathrm{SSM}})share of Var(log⁡10 util.)\mathrm{Var}(\log_{10}\,\mathrm{util.})
memory ρ/(1−ρ)2\rho/(1-\rho)^{2}6.22.0693%102%
dimension N DN\,D1.20.357%12%
readout ∥WC∥2\|W_{C}\|_{2}1.50.316%6%
input KBK_{B}1.80.283%3%
step KΔK_{\Delta}0.80.184%6%
transition KΔ+D ∥A∥2K_{\Delta}+\sqrt{D}\,\|A\|_{2}2.00.50-13%-11%
Appendix
Table 22: What drives the bound across the bank. log⁡10GSSM\log_{10}\mathcal{G}_{\mathrm{SSM}} is the sum of the logarithms of the factors of its recurrent-history term (plus a residual for the remaining terms), so the covariance of each factor with log⁡10GSSM\log_{10}\mathcal{G}_{\mathrm{SSM}} divided by its variance is an additive attribution that sums to one over all 1,296 parameter points. The last column attributes the variation of the utilization ∥∇L∥F/GSSM\|\nabla\mathcal{L}\|_{F}/\mathcal{G}_{\mathrm{SSM}} in the same way, with the measured gradient as an additional term. A negative share means the factor moves against the total: ∥A∥2\|A\|_{2} shrinks as ρ→1\rho\to 1 .
Figure 9: Additive attribution of the variance of ( 9(a) ) log⁡10GSSM\log_{10}\mathcal{G}_{\mathrm{SSM}} and ( 9(b) ) log⁡10\log_{10} utilization over all 1,296 parameter points to the factors of the bound (and, for the utilization, to the measured gradient). Bars sum to 100%, a negative bar moves against the total.
 β=0\beta=0β=0.7\beta=0.7β=1\beta=1
AlgorithmTrainValLast-10TrainValLast-10TrainValLast-10
FedAvg75.4 ±\pm 0.563.7 ±\pm 0.345.4 ±\pm 0.299.3 ±\pm 0.266.7 ±\pm 0.645.8 ±\pm 0.3143.5 ±\pm 0.7136.5 ±\pm 3.860.7 ±\pm 0.5
FedProx69.0 ±\pm 0.257.6 ±\pm 0.336.9 ±\pm 0.488.0 ±\pm 0.460.4 ±\pm 0.436.5 ±\pm 0.2119.3 ±\pm 0.4121.3 ±\pm 2.940.2 ±\pm 0.3
SCAFFOLD78.3 ±\pm 0.566.3 ±\pm 0.246.9 ±\pm 0.1102.4 ±\pm 0.467.3 ±\pm 0.345.2 ±\pm 0.1149.0 ±\pm 0.5131.6 ±\pm 3.649.7 ±\pm 0.2
FedAdam88.4 ±\pm 0.4164.8 ±\pm 1.653.0 ±\pm 0.3112.7 ±\pm 0.4160.0 ±\pm 2.150.8 ±\pm 0.1157.2 ±\pm 0.6355.2 ±\pm 13.861.7 ±\pm 0.3
FedAlign75.3 ±\pm 0.463.6 ±\pm 0.245.2 ±\pm 0.099.2 ±\pm 0.466.6 ±\pm 0.345.7 ±\pm 0.2143.5 ±\pm 0.5135.3 ±\pm 5.460.1 ±\pm 0.4
Appendix
Table 23: Effect of heterogeneity. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms on Mamba2 at every β\beta , averaged over the 2424 configurations.
 E=64E=64E=128E=128E=256E=256
AlgorithmTrainValLast-10TrainValLast-10TrainValLast-10
FedAvg151.2 ±\pm 0.5108.5 ±\pm 1.763.9 ±\pm 0.2100.9 ±\pm 0.585.5 ±\pm 1.148.7 ±\pm 0.266.1 ±\pm 0.372.9 ±\pm 2.039.2 ±\pm 0.4
FedProx137.7 ±\pm 0.3106.6 ±\pm 1.748.6 ±\pm 0.285.1 ±\pm 0.474.9 ±\pm 0.935.6 ±\pm 0.453.6 ±\pm 0.357.8 ±\pm 1.029.4 ±\pm 0.2
SCAFFOLD150.8 ±\pm 0.4105.6 ±\pm 1.761.0 ±\pm 0.3104.5 ±\pm 0.385.3 ±\pm 0.745.4 ±\pm 0.374.5 ±\pm 0.474.3 ±\pm 2.035.4 ±\pm 0.0
FedAdam167.3 ±\pm 0.5261.5 ±\pm 5.669.6 ±\pm 0.5113.0 ±\pm 0.4217.6 ±\pm 6.653.1 ±\pm 0.178.1 ±\pm 0.3200.9 ±\pm 5.242.8 ±\pm 0.1
FedAlign151.2 ±\pm 0.5108.3 ±\pm 2.163.6 ±\pm 0.2100.8 ±\pm 0.485.1 ±\pm 1.748.4 ±\pm 0.366.0 ±\pm 0.372.0 ±\pm 2.339.0 ±\pm 0.2
Appendix
Table 24: Effect of local steps. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms on Mamba2 at every EE , averaged over the 2424 configurations.
 Setting 1Setting 2Setting 3Setting 4
AlgorithmTrainValLast-10TrainValLast-10TrainValLast-10TrainValLast-10
FedAvg45.1 ±\pm 0.253.2 ±\pm 1.525.2 ±\pm 0.3107.3 ±\pm 0.096.5 ±\pm 4.251.7 ±\pm 0.6140.6 ±\pm 1.091.4 ±\pm 3.356.9 ±\pm 0.1131.3 ±\pm 0.5114.8 ±\pm 1.168.7 ±\pm 0.8
FedProx42.3 ±\pm 0.246.4 ±\pm 0.219.0 ±\pm 0.293.8 ±\pm 0.386.4 ±\pm 0.338.9 ±\pm 0.3117.7 ±\pm 0.687.3 ±\pm 3.143.0 ±\pm 0.3114.6 ±\pm 0.598.9 ±\pm 1.750.5 ±\pm 0.9
SCAFFOLD45.9 ±\pm 0.251.6 ±\pm 1.522.6 ±\pm 0.4112.0 ±\pm 0.095.6 ±\pm 3.447.5 ±\pm 0.2146.5 ±\pm 0.793.1 ±\pm 3.255.6 ±\pm 0.1135.2 ±\pm 0.5113.3 ±\pm 1.363.3 ±\pm 0.7
FedAdam51.4 ±\pm 0.1169.0 ±\pm 4.327.2 ±\pm 0.2118.2 ±\pm 0.4260.3 ±\pm 6.654.2 ±\pm 0.2158.0 ±\pm 1.1237.0 ±\pm 6.462.2 ±\pm 0.2150.2 ±\pm 0.2240.3 ±\pm 7.677.1 ±\pm 1.3
FedAlign45.0 ±\pm 0.050.8 ±\pm 1.624.8 ±\pm 0.2107.2 ±\pm 0.396.7 ±\pm 3.451.4 ±\pm 0.3140.6 ±\pm 0.991.8 ±\pm 2.756.7 ±\pm 0.1131.3 ±\pm 0.6114.6 ±\pm 1.268.5 ±\pm 0.9
Appendix
Table 25: Effect of the domain assignment. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms on Mamba2 for every domain setting (Setting 1: arxiv , github , stackexchange , wikipedia ; Setting 2: arxiv , c4 , stackexchange , wikipedia ; Setting 3: c4 , common_crawl , github , wikipedia ; Setting 4: arxiv , common_crawl , github , stackexchange ), averaged over the 1818 configurations.
 Depth 44 ( 14.714.7 M)Depth 1616 ( 20.320.3 M)
AlgorithmTrainValLast-10TrainValLast-10
FedAvg109.8 ±\pm 0.588.8 ±\pm 0.554.8 ±\pm 0.3102.3 ±\pm 0.489.1 ±\pm 2.646.4 ±\pm 0.2
FedProx93.2 ±\pm 0.479.0 ±\pm 1.339.5 ±\pm 0.491.0 ±\pm 0.380.5 ±\pm 1.236.2 ±\pm 0.1
SCAFFOLD112.8 ±\pm 0.488.0 ±\pm 1.450.1 ±\pm 0.0107.0 ±\pm 0.488.8 ±\pm 2.044.5 ±\pm 0.2
FedAdam122.6 ±\pm 0.3235.1 ±\pm 4.159.8 ±\pm 0.3116.2 ±\pm 0.5218.2 ±\pm 7.450.5 ±\pm 0.3
FedAlign109.7 ±\pm 0.589.1 ±\pm 1.954.4 ±\pm 0.2102.4 ±\pm 0.487.9 ±\pm 2.246.3 ±\pm 0.3
Appendix
Table 26: Effect of model depth. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms for the depth- 44 and depth- 1616 Mamba2 models, averaged over the 3636 configurations.
 K=4K=4K=8K=8K=16K=16K=32K=32
AlgorithmTrainValLast-10TrainValLast-10TrainValLast-10TrainValLast-10
FedAvg144.4 ±\pm 0.389.7 ±\pm 0.964.2 ±\pm 0.4141.4 ±\pm 0.487.3 ±\pm 0.561.6 ±\pm 0.2136.9 ±\pm 0.686.3 ±\pm 0.760.9 ±\pm 0.3129.0 ±\pm 0.386.2 ±\pm 0.661.1 ±\pm 0.2
FedProx131.6 ±\pm 0.282.4 ±\pm 0.849.0 ±\pm 0.2129.1 ±\pm 0.580.6 ±\pm 0.347.3 ±\pm 0.3125.2 ±\pm 0.479.9 ±\pm 0.546.8 ±\pm 0.2118.5 ±\pm 0.380.5 ±\pm 0.447.6 ±\pm 0.1
SCAFFOLD143.7 ±\pm 0.287.6 ±\pm 0.962.4 ±\pm 0.4139.5 ±\pm 0.583.8 ±\pm 0.558.7 ±\pm 0.2134.2 ±\pm 0.781.2 ±\pm 0.655.7 ±\pm 0.4126.8 ±\pm 0.479.7 ±\pm 0.553.0 ±\pm 0.2
FedAdam159.9 ±\pm 0.1187.1 ±\pm 2.069.5 ±\pm 0.6156.2 ±\pm 0.6183.0 ±\pm 1.267.0 ±\pm 0.4151.3 ±\pm 0.5181.7 ±\pm 2.066.1 ±\pm 0.4143.7 ±\pm 0.2182.8 ±\pm 1.766.5 ±\pm 0.5
FedAlign144.4 ±\pm 0.289.7 ±\pm 0.964.2 ±\pm 0.4141.4 ±\pm 0.487.3 ±\pm 0.561.6 ±\pm 0.2136.9 ±\pm 0.686.3 ±\pm 0.760.9 ±\pm 0.3129.0 ±\pm 0.386.2 ±\pm 0.661.1 ±\pm 0.2
Appendix
Table 27: Client-count ablation. Training, validation, and last-10-round validation perplexity of the nine algorithms with K∈{4,8,16,32}K\in\{4,8,16,32\} clients at β=0.7\beta=0.7 , depth 44 , E=64E=64 , averaged over the four domain settings.
 mamba2hybrid75hybrid50hybrid50althybrid25attn0
 SSM 1.00SSM 0.75SSM 0.50SSM 0.50SSM 0.25SSM 0.00
AlgorithmValLast-10ValLast-10ValLast-10ValLast-10ValLast-10ValLast-10
FedAvg43.4 ±\pm 0.430.7 ±\pm 0.341.5 ±\pm 0.826.5 ±\pm 0.543.3 ±\pm 0.629.3 ±\pm 0.452.6 ±\pm 1.231.3 ±\pm 0.254.4 ±\pm 1.232.7 ±\pm 0.769.6 ±\pm 1.142.3 ±\pm 1.1
FedProx41.9 ±\pm 0.425.4 ±\pm 0.239.8 ±\pm 0.421.0 ±\pm 0.241.4 ±\pm 0.323.6 ±\pm 0.250.4 ±\pm 0.724.7 ±\pm 0.452.2 ±\pm 1.326.8 ±\pm 0.967.7 ±\pm 1.034.8 ±\pm 0.9
SCAFFOLD42.2 ±\pm 0.429.4 ±\pm 0.339.2 ±\pm 1.022.5 ±\pm 0.541.3 ±\pm 0.725.4 ±\pm 0.850.3 ±\pm 1.227.1 ±\pm 0.352.5 ±\pm 1.229.6 ±\pm 1.368.3 ±\pm 1.041.8 ±\pm 0.9
FedAdam116.0 ±\pm 1.133.2 ±\pm 0.2116.5 ±\pm 3.430.7 ±\pm 0.5120.4 ±\pm 2.832.5 ±\pm 0.5141.7 ±\pm 3.938.8 ±\pm 0.2144.4 ±\pm 2.340.3 ±\pm 0.6180.8 ±\pm 4.659.4 ±\pm 1.1
Appendix
Table 28: Backbone ablation. Validation and last-10-round validation perplexity of the nine algorithms on the six depth- 44 backbones of Table 10 (Setting 1, β=0.7\beta=0.7 , E=64E=64 , K=4K=4 , the second header row gives the fraction of Mamba2 layers).
CheckpointModel11 epoch22 epochs44 epochsSelected λ\lambda / (λ,d)(\lambda,d)
Mamba2-130MZero-shot (pretrained)25.7 ±\pm 0.0–
Specialists (own task)81.8 ±\pm 1.888.4 ±\pm 0.589.2 ±\pm 0.2–
Specialists (other tasks)23.0 ±\pm 0.223.6 ±\pm 0.123.4 ±\pm 0.3–
Weight averaging34.4 ±\pm 2.136.3 ±\pm 1.737.4 ±\pm 0.6–
Task Arithmetic53.5 ±\pm 4.661.9 ±\pm 5.161.3 ±\pm 3.811 / 11 / 0.70.7
TIES-Merging47.0 ±\pm 2.055.0 ±\pm 3.956.6 ±\pm 3.6(1.3,0.1)(1.3,0.1) / (1.3,0.1)(1.3,0.1) / (1.3,0.1)(1.3,0.1)
Appendix
Table 29: One-shot merging of Mamba2 task specialists: task-averaged accuracy (%). For each public checkpoint, the accuracy averaged over AG News, DBpedia, TREC, and IMDB of the pretrained model, of the four specialists on their own task and on the other three tasks, and of the merged model produced by weight averaging, Task Arithmetic, and TIES-Merging from specialists fine-tuned for 11 , 22 , or 44 epochs.
 Mamba2-130MMamba2-370MMamba2-780M
ModelAG NewsDBpediaTRECIMDBAG NewsDBpediaTRECIMDBAG NewsDBpediaTRECIMDB
Zero-shot (pretrained)25.0 ±\pm 0.09.9 ±\pm 0.016.8 ±\pm 0.051.1 ±\pm 0.224.3 ±\pm 0.08.1 ±\pm 0.013.4 ±\pm 0.251.1 ±\pm 0.125.3 ±\pm 0.08.7 ±\pm 0.015.1 ±\pm 0.149.9 ±\pm 0.1
Specialists (own task)89.5 ±\pm 0.396.4 ±\pm 0.286.7 ±\pm 0.683.9 ±\pm 0.490.8 ±\pm 0.297.6 ±\pm 0.288.7 ±\pm 0.586.0 ±\pm 0.190.8 ±\pm 0.197.9 ±\pm 0.288.4 ±\pm 1.086.1 ±\pm 0.5
Weight averaging37.9 ±\pm 2.621.6 ±\pm 3.934.9 ±\pm 5.355.3 ±\pm 0.742.4 ±\pm 2.336.8 ±\pm 0.963.5 ±\pm 4.459.1 ±\pm 5.444.4 ±\pm 0.641.1 ±\pm 3.354.5 ±\pm 1.666.8 ±\pm 3.2
Task Arithmetic64.7 ±\pm 8.553.3 ±\pm 13.861.7 ±\pm 7.365.5 ±\pm 2.675.3 ±\pm 7.074.4 ±\pm 2.865.1 ±\pm 2.872.5 ±\pm 2.280.6 ±\pm 2.365.9 ±\pm 7.957.9 ±\pm 7.781.1 ±\pm 1.7
TIES-Merging72.2 ±\pm 8.451.7 ±\pm 17.946.2 ±\pm 13.156.2 ±\pm 0.177.3 ±\pm 6.873.4 ±\pm 3.662.9 ±\pm 3.069.9 ±\pm 0.883.4 ±\pm 1.469.5 ±\pm 8.252.4 ±\pm 11.683.1 ±\pm 1.1
Appendix
Table 30: Per-task accuracy (%) of the merged models after 4 specialist epochs. Columns are the evaluation tasks for each checkpoint, rows are the pretrained model, the specialist fine-tuned on that task, and the three merges.
CheckpointEpochsλ=0.2\lambda=0.2λ=0.3\lambda=0.3λ=0.5\lambda=0.5λ=0.7\lambda=0.7λ=1\lambda=1λ=1.3\lambda=1.3
Mamba2-130M1131.2 ±\pm 1.238.1 ±\pm 3.248.3 ±\pm 3.352.4 ±\pm 4.253.5 ±\pm 4.651.7 ±\pm 3.8
2232.6 ±\pm 1.141.3 ±\pm 1.955.5 ±\pm 2.061.9 ±\pm 4.161.9 ±\pm 5.155.8 ±\pm 5.4
4432.8 ±\pm 0.842.7 ±\pm 0.855.8 ±\pm 2.161.3 ±\pm 3.857.2 ±\pm 6.249.3 ±\pm 7.0
Mamba2-370M1140.2 ±\pm 1.949.5 ±\pm 0.659.4 ±\pm 2.961.1 ±\pm 2.558.6 ±\pm 3.454.1 ±\pm 5.6
2243.5 ±\pm 1.554.7 ±\pm 0.467.6 ±\pm 1.071.0 ±\pm 1.768.2 ±\pm 2.259.6 ±\pm 2.7
4444.5 ±\pm 1.555.4 ±\pm 0.568.6 ±\pm 0.571.8 ±\pm 1.266.3 ±\pm 2.554.0 ±\pm 3.3
Appendix
Table 31: Task Arithmetic scale sweep: task-averaged accuracy (%). θTA=θ0+λ∑mTm\theta_{\mathrm{TA}}=\theta^{0}+\lambda\sum_{m}\mathbf{T}_{m} over the four specialists, so λ=1/M=0.25\lambda=1/M=0.25 would be weight averaging and λ=1\lambda=1 adds every task vector at full length. The 780M sweep evaluated only λ∈{0.7,1.0}\lambda\in\{0.7,1.0\} .
Checkpointλ\lambdad=0.1d=0.1d=0.2d=0.2d=0.3d=0.3d=0.5d=0.5
Mamba2-130M0.50.534.7 ±\pm 1.533.7 ±\pm 0.932.6 ±\pm 0.832.1 ±\pm 0.8
0.80.846.1 ±\pm 1.344.0 ±\pm 1.142.2 ±\pm 1.141.2 ±\pm 1.4
1151.7 ±\pm 1.649.6 ±\pm 1.747.9 ±\pm 1.946.9 ±\pm 1.3
1.31.356.6 ±\pm 3.654.7 ±\pm 2.653.7 ±\pm 2.352.9 ±\pm 1.8
Mamba2-370M0.50.545.4 ±\pm 0.646.4 ±\pm 0.345.0 ±\pm 1.542.0 ±\pm 1.3
0.80.863.0 ±\pm 1.762.6 ±\pm 0.860.9 ±\pm 1.157.5 ±\pm 1.3
Appendix
Table 32: TIES-Merging sweep after 4 specialist epochs: task-averaged accuracy (%). Rows are the output scale λ\lambda applied to the task vector, columns the fraction dd of largest-magnitude entries kept per task vector. The 780M sweep evaluated only λ=1.3\lambda=1.3 with d∈{0.1,0.5}d\in\{0.1,0.5\} in addition to the default (λ,d)=(1.0,0.2)(\lambda,d)=(1.0,0.2) .
 Specialists (own task)Weight averaging
μ\mu11 epoch22 epochs44 epochs11 epoch22 epochs44 epochs
00 (reference)81.8 ±\pm 1.888.4 ±\pm 0.589.2 ±\pm 0.234.4 ±\pm 2.136.3 ±\pm 1.737.4 ±\pm 0.6
10−510^{-5}82.6 ±\pm 3.088.1 ±\pm 0.988.9 ±\pm 0.633.3 ±\pm 0.935.7 ±\pm 0.637.3 ±\pm 0.6
10−410^{-4}82.5 ±\pm 3.587.6 ±\pm 0.388.1 ±\pm 0.332.5 ±\pm 2.033.1 ±\pm 1.134.0 ±\pm 0.8
10−310^{-3}80.6 ±\pm 2.685.3 ±\pm 0.686.1 ±\pm 0.130.4 ±\pm 1.130.3 ±\pm 0.530.6 ±\pm 0.5
10−210^{-2}75.3 ±\pm 3.179.7 ±\pm 0.981.9 ±\pm 0.427.1 ±\pm 0.727.6 ±\pm 0.128.0 ±\pm 0.1
Appendix
Table 33: Weight averaging of proximally fine-tuned specialists (Mamba2-130M): task-averaged accuracy (%). Each specialist is fine-tuned with the FedProx local objective Lm(θ)+μ2∥θ−θ0∥F2\mathcal{L}_{m}(\theta)+\frac{\mu}{2}\|\theta-\theta^{0}\|_{F}^{2} for 11 , 22 , or 44 epochs and the four specialists are merged by weight averaging

Explore similar work

CardsList
Jul 31, 2026cs.LG

Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients

Federated learning of foundation models faces a fundamental resource-asymmetry challenge: the institutions holding the most valuable domain-specific data cannot host billion-parameter models. Existing heterogeneous federated approaches attempt to bridge this gap through parameter-efficient tuning, model pruning, or knowledge distillation, yet each trades away a critical property, whether full-model memory reduction, architectural self-containedness, or representational fidelity, leaving the core tension unresolved. We propose FedSLM, a parameter-centric framework for federated fine-tuning with heterogeneous compressed clients. FedSLM uses SVD-based decomposition to produce self-contained client models, whose low-rank subspaces form nested manifolds that are structurally compatible for aggregation. It then applies a two-stage protocol that synchronizes lightweight adapters within compression groups and fuses full-rank reconstructions across groups via structural alignment. Finally, a weak-to-strong elicitation step with auxiliary confidence loss transfers the aggregated knowledge to the full-scale server, while an explicit bias--variance trade-off mitigates compression artifacts. We provide theoretical guarantees for adapter-level aggregation, subspace-alignment bounds for cross-group fusion, and a characterization of how the confidence loss mitigates weak-supervision noise. Experiments on natural language and vision--language benchmarks show that FedSLM outperforms existing federated baselines under both IID and non-IID partitions, while client models operate at roughly 50% of the GPU memory required by the full model.

Shengkun Zhu, Jinshan Zeng, Zhihua Allen-Zhao +5

The Hong Kong Polytechnic University · OceanBase, Ant Group · School of Management, Xi’an Jiaotong University +3

FedavgParameter-Efficient Fine-Tuning MethodsWireless Foundation ModelsModel Fine-TuningFoundation ModelClients
AskPDF
Apr 28, 2026cs.LG

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

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 heterogeneity-correction 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-variance assumptions, SSF attains a non-asymptotic rate of order O~(1/T+1/NKT)\widetilde{\mathcal{O}}(1/T+1/\sqrt{NKT}). Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.

Shuchen Zhu, Zhengyang Huang, Yuqi Xu +1

Peking University · Beihang University

Heterogeneous Federated LearningFederated LearningHeterogeneitySubspace
AskPDF
Jun 21, 2026cs.LG

Subspace-Constrained Federated Learning with Low-Rank Adaptation

Federated low-rank adaptation methods are attractive for fine-tuning large models under communication and privacy constraints, but heterogeneous client data can induce geometric misalignment between local low-rank updates. We study whether this subspace misalignment leads to destructive aggregation and slower convergence in LoRA-based federated learning. We propose a subspace-regularized federated LoRA objective that encourages local client updates to remain close to a shared global reference subspace. We present a complete empirical evaluation on two pretrained models, RoBERTa-large and SmolLM-360M, over HellaSwag in a non-IID 10-client federated setting, across 3 random seeds (42, 43, 44), yielding 24 total experimental runs (4 methods x 3 seeds x 2 models). On RoBERTa-large, Subspace-Reg achieves the strongest mean best accuracy (0.454 +/- 0.023), mean final accuracy (0.429 +/- 0.011), and lowest final loss (1.363) across all three seeds, outperforming FedAvg, SVD redistribution, and FedSVD baselines by a large margin. On SmolLM-360M, FedAvg leads on accuracy, revealing that accuracy gains are model-dependent. Crucially, Subspace-Reg achieves near-perfect basis overlap, approximately 0.9999, on both models and across all seeds, versus 0.958 to 0.991 for all baselines, providing robust support for the geometric alignment hypothesis. The code is publicly available at https://github.com/sadia-sigma-lab/Subspace-Constrained-Federated-learning-with-Lora.

Neranjan Senarath, Rohit Muralitharan, Sadia Asif

Department of ECSE Rensselaer Polytechnic Institute Troy, NY 12180, USA · Department of Computer Science Rensselaer Polytechnic Institute Troy, NY 12180, USA

Federated LearningSubspace
AskPDF