PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.02659v1cs.LGcs.AImath.OC

TL;DR

Federated convergence behavior for modern selective SSMs remains poorly characterized, while standard FL methods are largely architecture-agnostic. The paper derives architecture-aware bounds and tests their relevance, finding no gradient-bound violations and algorithm rankings that correlate highly with validation performance.

  • Problem

    Federated-learning methods are largely architecture-agnostic, and convergence behavior for modern selective SSMs remains poorly characterized despite their distinctive stability and parameterization.

    Source: §1 · pp. 1–2

  • Method

    The paper derives selective-SSM gradient and smoothness bounds and uses them to analyze federated convergence in terms of architecture-specific quantities.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    Theory-motivated algorithm rankings highly correlated with validation performance, and teacher-sequence experiments detected no gradient-bound violations.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The ranking agreement supports the practical relevance of SSM-specific quantities for interpreting federated-algorithm performance.

    Source: §1 · p. 2

    The experiments found the derived bounds conservative in absolute value, with looseness associated with parameter-envelope growth and saturation of the sequence-level total-variation metric.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

SSMs offer linear-time sequence modeling as an alternative to transformers’ quadratic self-attention cost, but convergence behavior for modern selective SSMs in federated settings remains largely uncharacterized. The paper derives architecture-aware bounds and examines their relevance through teacher-generated sequences and federated language-modeling experiments.

1 INTRODUCTION

  • SSMs avoid transformers’ quadratic self-attention cost with linear-time sequence processing and constant-sized latent state.The passage gives complexities O(L) and O(1), respectively.Source: §1 · p. 1
  • Federated convergence for modern selective SSMs remains poorly characterized because most FL methods are architecture-agnostic and prior SSM studies address non-selective or specialized personalized settings.The paper asks how selective SSM structure can support architecture-aware convergence bounds for distributed learning.Source: §1 · p. 2
  • The paper derives selective-SSM gradient bounds and FedAvg convergence rates that expose dependencies on SSM parameter blocks and block-wise spectral-norm envelopes.The authors describe this as the first architecture-aware FL convergence rate for selective SSMs, to their knowledge.Source: §1 · p. 2
  • Teacher-sequence tests found no gradient-bound violations, while theory-motivated algorithm rankings highly correlated with validation performance on heterogeneous-domain datasets.The experiments also observed greater client disagreement as ρ approaches one and greater drift with more local steps.Source: §1 · p. 2

2 RELATED WORK

Related work spans structured and selective SSMs, federated optimization, and specialized SSM applications. The paper identifies weight-space aggregation of independently trained selective SSMs as an unresolved distributed-learning question.

2 RELATED WORK

  • Modern SSM work developed structured parameterizations and selective mechanisms for capturing long-range dependencies, with applications across language modeling, vision, and sequence modeling.The passage also describes Structured State Space Duality as connecting SSMs with linear attention.Source: §2 · p. 2
  • Federated methods address statistical heterogeneity, communication costs, and client drift, while SSM-specific FL studies have examined system identification and personalized meta-learning.FedAlign targets differing parameter basins in system identification; FedSSM uses SSMs as a meta-learning tool for personalized FL.Source: §2 · p. 2
  • Prior hidden-state, context, or representation merging differs from weight-space aggregation of independently trained selective SSMs, whose distributed optimization guarantees remain open.This distinction motivates connecting selective-SSM architecture to distributed-training behavior.Source: §2 · p. 3

3 SYSTEM MODEL

The system model combines a selective state-space recurrence with decentralized training across clients holding local sequence data. Input-dependent parameters govern token processing, while local optimization and server aggregation repeatedly update the global model.

3 SYSTEM MODEL

  • The global SSM parameter tuple comprises input, output, discretization, and transition parameters, with the transition parameter representing diagonal entries of a structured matrix.The latent state evolves over time to map continuous inputs to outputs.Source: §3 · p. 3
  • The model uses first-order exponential-Euler discretization, retaining the exponential state transition while approximating the input contribution to the discrete recurrence.The recurrence maps a length-S input sequence to outputs through compressed latent states.Source: §3 · p. 3
  • Selective SSM parameters depend on the current token: the time scale controls discretization, the transition controls memory retention and decay, and input and output matrices modulate state updates and projections.The hidden state and output are computed from these dynamically generated quantities at each sequence step.Source: §3 · pp. 3–4
  • At each communication round, the server broadcasts global weights and clients perform E local SGD steps on their local objectives; FedProx adds a proximal penalty to mitigate divergence from data heterogeneity.FedAvg uses the expected local sequence loss, and the clients’ resulting models are aggregated into the next global model.Source: §3 · p. 4

4 THEORETICAL ANALYSIS

The analysis derives bounds from selective-SSM structure rather than assuming architecture-agnostic smoothness and disagreement constants. It connects stability, parameter sensitivity, and client displacement to gradient growth and federated convergence behavior.

4 THEORETICAL ANALYSIS

  • The convergence analysis tracks how recurrent stability, input-dependent time scales, and projection norms control gradient growth, client drift, and aggregation error.It establishes guarantees for FedAvg and FedProx alongside a single-layer gradient bound.Source: §4 · p. 5

4.1 GRADIENT NORM BOUND FOR SELECTIVE SSMS

  • Under stability and regularity conditions, the single-layer gradient bound is independent of sequence length and explicitly depends on ρ and the core selective-SSM parameter blocks.Bounding hidden states and their parameter Jacobians controls the recurrent contributions through geometric sums when ρ = exp(−γδ) < 1.Source: §4.1 · pp. 5–6

4.2 CONVERGENCE ANALYSIS

  • The FedAvg guarantee derives smoothness, stochastic variance, and client disagreement from the recurrence, making its model dependence explicit in parameter envelopes and ρ.The analysis does not require separately assumed model smoothness or gradient-dissimilarity constants.Source: §4.2 · pp. 6–7
  • Gradient and smoothness estimates have no explicit sequence-length factor but worsen as ρ approaches one, with stronger stability penalties in local-drift terms.The local-drift term contains M^6_R/(1 − ρ)^12 versus M^4_R/(1 − ρ)^8 for the baseline stochastic term; the heterogeneity contribution vanishes at E = 1 and does not decrease with batch size when E > 1.Source: §4.2 · p. 7
  • Greater client displacement enlarges parameter envelopes, while FedProx trades smaller envelopes from restricted movement against increased local smoothness.The trajectories must remain within a neighborhood contained in the stable parameter set.Source: §4.2 · p. 7

5 EXPERIMENTS

Experiments test theoretical bounds on teacher-generated sequences and compare nine federated algorithms on Mamba2; observed algorithm performance directions tracked validation behavior.

Validation of Derived Bounds

  • Frozen one-layer teachers generated 72 datasets; learners trained on 4,096 sequences, and 1,411,776 gradient, smoothness, and bound-variation checks found no violations.The 103,200 federated checks covered 1,200 FedAvg/FedProx runs, including step sizes beyond the theorems’ range.Source: §5 · pp. 7–8, Table 1 · p. 8

Memory amplifies gradient and curvature sensitivity

  • Measured gradient and Hessian norms grew approximately as (1−ρ)^−0.5 and (1−ρ)^−1.0, while their bounds grew as (1−ρ)^−1.8 and (1−ρ)^−3.6.The fitted curvature exponent was twice the gradient exponent for both measured quantities and certified constants.Source: Figure 2 · p. 8, §5 · p. 8

Recurrent sensitivity and block-wise gradients

  • The recurrent-Jacobian term accounted for at least 99% of GSSM; measured gradient energy shifted from WB, WC at ρ=0.3 to A at ρ=0.999.The corresponding measured shares were 86% for WB, WC and 85% for A; fitted bound exponents varied with vocabulary and state dimensions.Source: §5 · p. 8

Validation of Derived Bounds

  • Observed client drift grew approximately as E^0.87, while client disagreement increased as ρ approached one, qualitatively matching the FedAvg analysis.The drift trend is consistent with the linear accumulation bound ηEGSSM,R.Source: §5 · p. 8

Results and Discussion

  • Expected performance directions correlated with observed log-perplexity changes at 0.71 overall and 0.81 for β=1; 13/16 outcomes matched across β∈{0,1}.SCAFFOLD improved over FedAvg in 23/24 configurations at β=1 and deteriorated in 22/24 at β=0.Source: Table 2 · p. 9, §5 · pp. 9–10

Limitations and outlook

  • The bounds were conservative in absolute value, and radius-envelope growth and total-variation saturation contributed to looseness; token-level dissimilarity is proposed for partition-sensitive heterogeneity.The authors also propose data-dependent constants and envelopes that track actual local trajectories.Source: §5 · p. 9

Practical Federated Optimization

  • Mamba2 experiments compared nine algorithms across 72 configurations, three seeds, six text domains, four domain assignments, and local-step and heterogeneity settings.Models had depth 4 or 16, trained across four clients for 50 rounds, and were evaluated using mean last-ten-round validation perplexity.Source: §5 · pp. 7, 9

Theory-Guided Expectations

  • FedProx’s proximal term bounds drift but adds a smoothness cost; SCAFFOLD and FedDyn may reduce corrected disagreement when control states track client-gradient differences.FedExP, FedAdam, and DiLoCo-style methods modify server updates, while FedAlign and small-perturbation FedSAM may behave like FedAvg without a shared convergence guarantee.Source: §5 · p. 9

6 CONCLUSION

The paper derives and validates architecture-aware bounds for distributed learning with selective SSMs, then uses them to interpret federated optimization experiments.

6 CONCLUSION

  • The study connects federated optimization to recurrent stability, input-dependent discretization, projection norms, client drift, and heterogeneity, and reports validation and practical experiments.The conclusion describes the derived bounds as relevant to distributed learning with selective SSMs.Source: §6 · p. 10
  • Generative AI assisted with proof feedback, hypotheses, methods, experiments, code, and manuscript presentation; authors reviewed the work and take responsibility for the final content.The authors state that AI did not develop theoretical models, generate synthetic datasets, or provide critical ingredients for proofs.Source: AI USE STATEMENT · p. 10
  • The authors provide complete theoretical proofs and document empirical setups, tuning procedures, hardware, code, and experiment-running instructions.These reproducibility materials are provided in Appendix G and supplementary material.Source: REPRODUCIBILITY STATEMENT · p. 10

APPENDIX

A SELECTIVE SSM GRADIENT BOUND PROOF

This section proves uniform single-layer selective SSM gradient bounds and develops smoothness, stochastic-gradient, and client-heterogeneity controls to support convergence analysis of FedAvg and FedProx, including a FedAvg bound under stated stability and regularity conditions.

A SELECTIVE SSM GRADIENT BOUND PROOF

  • The single-layer gradient bound holds uniformly for every admissible sequence length and pair at a fixed parameter tuple, and also bounds the averaged empirical loss.The proof derives the loss-gradient bound from the token-level bound; its right-hand side is independent of sequence length and sample.Source: A SELECTIVE SSM GRADIENT BOUND PROOF · p. 21

B CONVERGENCE OF FEDERATED SSM BASELINES

  • The federated-baselines section develops lemmas to instantiate standard convergence results with selective-SSM parameters, then addresses FedAvg and FedProx bounds.It introduces the proof of Theorem B.5 and states the corresponding FedProx theorem.Source: B CONVERGENCE OF FEDERATED SSM BASELINES · p. 21

B.1 AUXILIARY LEMMAS FOR SSM FEDERATED CONVERGENCE

  • The auxiliary lemmas replace generic convergence quantities with selective-SSM-specific controls, including gradient, smoothness, second-moment, stochastic-variance, and client-heterogeneity bounds.These lemmas are used to state convergence bounds for FedAvg and FedProx.Source: §B.1 · p. 21

B.1.1 LOCAL LIPSCHITZ CONTINUITY OF THE SSM GRADIENT BOUND LEMMA AND PROOF

  • On a finite-radius neighborhood within the admissible parameter set, the gradient-bound function GSSM is Lipschitz continuous with a finite constant.The constant depends on the neighborhood radius and reference parameter-block norms, with model dimensions and stated assumptions fixed.Source: §B.1.1 · pp. 22–23

B.1.2 LOCAL SSM SMOOTHNESS LEMMA PROOF

  • The local SSM smoothness proof uses constants uniform over a parameter neighborhood, with estimates independent of sequence length, admissible sample, and token position.The section identifies LSSM,R as the local smoothness constant established in Lemma 4.7.Source: §B.1 · p. 21, §B.1.2 · p. 24
  • The proof establishes local loss smoothness by bounding recurrent-state, output, and token-loss derivatives within a stable model neighborhood.The smoothness inequality holds uniformly over admissible sequence pairs and finite sequence lengths.Source: §B.1.2 · pp. 27–28

B.1.3 SSM LOSS GRADIENT SECOND-MOMENT CONSTANT LEMMA AND PROOF

  • Under Assumptions 4.2 and 4.3, Lemma B.2 establishes a second-moment bound for sequence-level stochastic gradients in a stable local neighborhood.The bound is expressed using constants defined for the selective SSM’s local parameter neighborhood.Source: §B.1.3 · p. 29

B.1.4 SSM STOCHASTIC GRADIENT VARIANCE BOUND LEMMA AND PROOF

  • Lemma B.3 derives a mini-batch stochastic-gradient variance bound from unbiased, independent sequence samples.Its assumptions include those of Lemma B.2 and Assumption 4.6.Source: §B.1.4 · pp. 29–30

B.1.5 SSM CLIENT HETEROGENEITY BOUND UNDER TOTAL VARIATION LEMMA AND PROOF

Define the local client heterogeneity constant

  • The local client heterogeneity constant is obtained by applying the pointwise gradient bound and taking a supremum over the stable neighborhood.The proof aggregates the resulting client-wise bounds using the client weights.Source: Define the local client heterogeneity… · p. 31

B.2 THEOREM B.5: FEDAVG SSM CONVERGENCE AND PROOF

  • Theorem B.5 gives a FedAvg convergence bound for a single-layer selective SSM using SSM-specific smoothness, gradient, variance, and heterogeneity bounds.The result assumes full client participation, a step size satisfying the theorem’s conditions, and common stable neighborhoods containing local iterates.Source: §B.2 · pp. 32–33

B.3 THEOREM B.6: FEDPROX SSM CONVERGENCE AND PROOF

The section derives a convergence bound for FedProx with selective SSMs and extends the corresponding gradient and convergence-bound derivations to finite-depth multi-layer SSMs.

B.3 THEOREM B.6: FEDPROX SSM CONVERGENCE AND PROOF

  • FedProx adds a proximal penalty around the broadcast model to each client's local objective, changing local optimization dynamics.The theorem considers full participation with constant local step size and proximal coefficient.Source: §B.3 · pp. 33–34
  • Theorem B.6 establishes a convergence bound for FedProx with selective SSMs.The theorem's stated bound holds under its specified setup and assumptions.Source: §B.3 · pp. 34, 35, 41
  • All local iterates and the aggregated model remain within the specified neighborhood of the broadcast model.That neighborhood is contained in the stable parameter set.Source: §B.3 · p. 35

C MULTI-LAYER SELECTIVE SSM PROOFS

C.1 MULTI-LAYER SELECTIVE SSM GRADIENT BOUND PROOF

The proof bounds gradients for sequential multi-layer selective SSMs under bounded layer inputs and layer-wise stability assumptions, uniformly over admissible sequence lengths and inputs for fixed model parameters. It also establishes layerwise parameter neighborhoods and envelopes for the federated-convergence proofs under uniform stability and regularity assumptions.

C.1 MULTI-LAYER SELECTIVE SSM GRADIENT BOUND PROOF

  • The analysis assumes uniformly bounded inputs to each layer and defines stability margins separately for every layer.The bounded-input constants are independent of sequence length, input-target pair, and token index.Source: §C.1 · p. 43
  • For stable multi-layer selective SSM parameters satisfying the stated assumptions, the theorem bounds the loss-gradient norm for every finite sequence length and admissible input-target pair.The model has finite depth Q, and the proof combines layer-wise gradient bounds.Source: §C.1 · p. 43
  • The proof applies the chain rule recursively, propagating each layer’s parameter Jacobian through input-output Jacobians of subsequent layers.Each earlier layer’s parameters affect later outputs through the inputs passed to those layers.Source: §C.1 · p. 46
  • Layer-wise gradient bounds hold uniformly over tokens and admissible input sequences, enabling their combination into a bound for the full model.For each layer, the proof reuses the single-layer argument while holding that layer’s input fixed with respect to its own parameters.Source: §C.1 · p. 47
  • Layerwise sensitivity bounds are combined through nested sums to obtain hidden-state gradient bounds for every layer and token position.The bounds hold for each layer and token position, with an empty product interpreted as one.Source: §C.1 · pp. 51–52
  • The loss-gradient bound holds uniformly over all admissible sequence lengths and input-target pairs for a fixed full parameter tuple.The bound is independent of sequence length and input-target pair, under the stated assumptions.Source: §C.1 · p. 52

C.2 AUXILIARY LEMMAS FOR MULTI-LAYER SSM FEDERATED CONVERGENCE

  • The auxiliary lemmas define layerwise parameter neighborhoods and envelopes, assuming uniform finite regularity constants and fixed stability margins.The envelopes bound parameter-dependent factors throughout the neighborhoods around the reference model.Source: §C.2 · pp. 52–53

C.2.1 LOCAL MULTI-LAYER SSM SMOOTHNESS LEMMA AND PROOF

This appendix section develops multi-layer SSM regularity bounds and supplementary convergence corollaries for selective-SSM FedAvg and FedProx and single-layer FedAvg, making dependence on stability, parameter norms, and depth explicit.

C.2.1 LOCAL MULTI-LAYER SSM SMOOTHNESS LEMMA AND PROOF

  • The local multi-layer loss is L(Q)-smooth, with a bound independent of sequence length and admissible sequence; client objectives and their weighted global objective inherit this smoothness.The proof bounds first- and second-order output sensitivities uniformly across layers and token positions, using finite neighborhood envelopes and recurrent stability ρ_q < 1.Source: §C.2.1 · pp. 54, 59

C.2.2 MULTI-LAYER SSM GRADIENT SECOND-MOMENT CONSTANT LEMMA AND PROOF

  • Lemma C.5 establishes a stochastic sequence-gradient second-moment bound within a stable local neighborhood around the round-r broadcast model.The lemma assumes Assumptions 4.2 and C.1; its constant uses the previously defined local neighborhood bounds.Source: §C.2.2 · p. 59

C.2.3 MULTI-LAYER SSM STOCHASTIC GRADIENT VARIANCE BOUND LEMMA AND PROOF

  • For independent sampled sequences, the mini-batch stochastic-gradient variance follows the standard variance identity and is bounded as stated in Lemma C.6.The lemma relies on the assumptions of Lemma C.5 and Assumption 4.6.Source: §C.2.3 · pp. 60–61

C.2.4 MULTI-LAYER SSM CLIENT HETEROGENEITY BOUND UNDER TOTAL VARIATION LEMMA AND PROOF

  • Lemma C.7 bounds multi-layer SSM client heterogeneity under total variation between each local data distribution and the global mixture.The result assumes Lemma C.5 and Assumption 4.6, with client weights summing to one.Source: §C.2.4 · p. 61

Define the local client heterogeneity constant

  • The local client heterogeneity constant is bounded by applying a total-variation inequality to the pointwise sequence-gradient bound, then taking the supremum over the local model neighborhood.The proof weights and sums the client-wise inequalities, yielding a bound independent of the model within that neighborhood.Source: Define the local client heterogeneity… · p. 62

If the step size satisfies

  • The multi-layer proofs transfer federated optimization arguments to the full parameter tuple, using common smoothness and noise bounds from the SSM lemmas and convex neighborhoods containing local iterates and aggregates.For FedAvg, conditional independence across clients supplies the weighted-noise variance bound; the proof substitutes the resulting constants under the stated step-size conditions.Source: If the step size satisfies · p. 65

C.3 COROLLARY C.8: FEDAVG MULTI-LAYER SSM CONVERGENCE AND PROOF

  • Corollary C.8 derives a multi-layer selective-SSM FedAvg convergence inequality under bounded neighborhoods, common smoothness and gradient bounds, conditional sampling, and specified step-size conditions.The proof applies the federated optimization argument for FedAvg to the full multi-layer parameter tuple.Source: §C.3 · p. 62, If the step size satisfies · p. 65

C.4 COROLLARY C.9: FEDPROX MULTI-LAYER SSM CONVERGENCE AND PROOF

  • Corollary C.9 gives a multi-layer selective-SSM FedProx convergence guarantee, with assumptions including a positive proximal coefficient and uniform architectural and loss bounds.Its proof bounds local iterates within reduced neighborhoods and applies the convergence argument using smoothness, variance, and client-disagreement bounds; the common proximal term cancels in gradient disagreement.Source: §C.4 · pp. 65–66, If the step size satisfies · p. 68

D SUPPLEMENTARY COROLLARIES

  • The supplementary corollaries make dependence on SSM stability, parameter norms, network depth, and related architectural properties explicit.They simplify results established in preceding sections and the main manuscript.Source: D SUPPLEMENTARY COROLLARIES · p. 69

D.1 SINGLE-LAYER CONVERGENCE RATE COROLLARIES

  • The single-layer corollaries provide convergence-rate guarantees for FedAvg and FedProx, retaining the assumptions of their referenced results.The section treats O expressions as upper bounds on nonnegative quantities.Source: §D.1 · p. 69

D.1.1 COROLLARY 4.8: SINGLE-LAYER FEDAVG CONVERGENCE RATE PROOF

  • The proof of Corollary 4.8 substitutes the bounds from Lemmas 4.7 and B.2 into Theorem B.5, collects the local-drift terms, and uses P2 ≤ 1.These steps yield the inequality stated in Corollary 4.8.Source: §D.1.1 · p. 69

D.1.2 COROLLARY D.1: SINGLE-LAYER FEDPROX CONVERGENCE RATE AND PROOF

These appendix results simplify selective-SSM gradient and convergence bounds under common layer envelopes and state when bounds depend on depth. They also establish a one-step setting in which FedProx and FedAvg updates coincide.

D.1.2 COROLLARY D.1: SINGLE-LAYER FEDPROX CONVERGENCE RATE AND PROOF

  • The single-layer FedProx bound has reduced envelopes but also extra μ terms, so it does not guarantee improvement for every positive proximal coefficient.Source: §D.1.2 · p. 69

D.2.1 COROLLARY D.2: MULTI-LAYER GRADIENT UNDER COMMON LAYER BOUNDS AND PROOF

  • With common envelopes fixed across depth, the gradient upper bound is uniform for κ < 1, at most √Q for κ = 1, and may grow exponentially for κ > 1.Temporal stability ρ < 1 alone does not imply κ < 1, and an exponentially growing upper bound does not establish exponential growth in actual gradients. Common layer bounds need not mean identical weights or stability factors.Source: §D.2.1 · p. 70, §D.2 · p. 70

D.2.2 COROLLARY D.3: MULTI-LAYER GRADIENT RATE UNDER COMMON LAYER BOUNDS AND PROOF

  • The multi-layer gradient-rate corollary’s implicit constants depend on N, D, Lx, and Lℓ, but not depth Q, sequence length, parameter norms, or stability factors.Source: §D.2.2 · p. 71

D.3.1 COROLLARY D.4: COMMON-LAYER SMOOTHNESS AND DEPTH DEPENDENCE AND PROOF

  • The multi-layer smoothness reduction replaces layerwise sensitivity recursions with three finite geometric sums; its depth statement varies Q while fixing common envelopes and regularity constants.Source: §D.3 · p. 71, §D.3.1 · p. 72

D.3.2 COROLLARY D.5: MULTI-LAYER FEDAVG UNDER COMMON LAYER BOUNDS AND PROOF

  • The common-layer FedAvg corollary applies bounds bGQ and bLQ under Corollary C.8 assumptions, with depth dependence inherited from the explicit common-bound formulas.Source: §D.3.2 · p. 73

D.3.3 COROLLARY D.6: MULTI-LAYER FEDPROX UNDER COMMON LAYER BOUNDS AND PROOF

  • The common-layer FedProx corollary retains original and effective radii while applying common bounds; its implicit constant is numerical and independent of μ.Source: §D.3.3 · pp. 73–74

D.4 SSM CONVERGENCE UNDER DISTRIBUTED LEARNING

  • With one local gradient step per exchange, FedProx’s update and weighted aggregate match FedAvg for every μ > 0, without an ημ restriction.The convergence bounds apply under the same FedAvg assumptions and step-size conditions.Source: §D.4 · p. 74

E DISCUSSION ON ADDITIONAL DISTRIBUTED LEARNING FRAMEWORKS

The discussion interprets additional federated algorithms through selective-SSM displacement, disagreement, and stability quantities. It also connects weighted averaging of independently adapted experts to one round of FedAvg, while noting limits on extending the bounds to broader merging settings.

E DISCUSSION ON ADDITIONAL DISTRIBUTED LEARNING FRAMEWORKS

  • SCAFFOLD changes effective heterogeneity rather than the SSM smoothness constant, but its correction may increase local drift and offer little benefit when data are nearly IID.Under non-IID data, the correction may reduce the statistical drift term, which is amplified by architecture-dependent gradient and smoothness quantities.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · pp. 75–76
  • FedAdam leaves local SSM drift unchanged and may enlarge the global update radius when its server step is aggressive or its moment estimates poorly aligned.Compared with FedProx and SCAFFOLD, it neither explicitly constrains local drift nor corrects client heterogeneity.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · p. 76
  • FedDyn combines proximal displacement damping with disagreement reduction when its centered states track client-gradient differences, but its benefits must outweigh increased smoothness and preserve stable containment.The discussion expects a well-tuned FedDyn to improve on FedAvg, especially under heterogeneous data.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · p. 77
  • FedExP rescales when client-update disagreement is sufficiently large relative to the average update, potentially requiring a larger stable neighborhood without correcting disagreement.The discussion expects scaling to be mostly inactive on IID data and possibly harmful under heterogeneity when displacement and stability costs outweigh acceleration.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · pp. 78–79
  • DiLoCo-style server momentum can enlarge the displacement bound when αo/(1 − βo) > 1, and accumulated or misaligned momentum may outweigh acceleration without correcting client disagreement.The discussion treats worse performance than FedAvg as a possibility, not a guaranteed outcome.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · p. 79
  • FedProx and FedDyn can damp client displacement at a smoothness cost, while SCAFFOLD and FedDyn can reduce corrected disagreement when their controls track client-gradient differences.The expected benefit is smaller when IID data leave less disagreement to correct.Source: E DISCUSSION ON ADDITIONAL DISTRIBUTED… · p. 79

F.1 MODEL MERGING WITH SSMS MOTIVATION

  • Weighted averaging of experts adapted from a shared initialization is exactly one round of FedAvg when objective and aggregation weights match.More general merging operators retain independent adaptation followed by aggregation, but their different combination rules limit the current theory.Source: §F.1 · p. 80

F.2 MODEL MERGING METHODS: CONNECTIONS TO SELECTIVE SSM ANALYSIS

  • Applying the federated bounds to one-shot merging may require a large neighborhood around the shared initialization, making bounds loose or stability and step-size conditions harder to satisfy.A common stable neighborhood and convex averaging do not by themselves guarantee predictive performance across contributing domains.Source: §F.1 · p. 80

F + LSSM,R

The merging analysis relates local loss predictions to update geometry and SSM smoothness, while the experimental setup covers synthetic validation, federated language modeling, and one-shot merging.

F + LSSM,R

  • The local loss bound after merging depends on update alignment with −∇f(W0), the merge operator, and the SSM smoothness over the chosen neighborhood.A larger radius can require larger parameter envelopes and smoothness constants, loosening the bound.Source: F + LSSM,R · pp. 80–81
  • Weighted parameter averaging keeps the merged model within the experts’ common stable neighborhood, preserving the same local smoothness bound.The weighted squared gradient disagreement is bounded by 4G2SSM,RHTV, jointly reflecting data heterogeneity and SSM sensitivity.Source: F + LSSM,R · p. 81
  • Task Arithmetic scales combined task vectors, but shrinking updates can suppress useful adaptation and cannot recover directions lost through cancellation.Larger scales require a separate containment check; the directional term scales with λ and the quadratic remainder with λ2.Source: F + LSSM,R · p. 81
  • TIES-Merging may leave the experts’ common neighborhood, and the bounds do not establish that it improves stability or reduces merged-model loss.Its raw-magnitude and sign choices do not account for interactions among A, W∆, WB, and WC through the SSM recurrence.Source: F + LSSM,R · p. 81
  • The analysis is conditional on uniform regularity in a common convex stable neighborhood and, for federated convergence, containment of full adaptation trajectories.Extensive fine-tuning may require large parameter envelopes and yield loose bounds.Source: F + LSSM,R · p. 81

G DETAILED EXPERIMENTAL SETUP

  • The detailed experiments cover synthetic validation of theoretical bounds, federated language modeling with selective SSMs on RedPajama, and supplementary one-shot merging with pretrained Large Mamba2 models.The setup details datasets, models, and training and validation procedures for reproducibility.Source: G DETAILED EXPERIMENTAL SETUP · pp. 81–82

G.1 DATASETS

This section describes the synthetic and text-based datasets, SSM models, and synthetic and practical federated evaluation protocols, including model-merging and reproducibility procedures.

G.1 DATASETS

  • Accepted teachers improved on uniform NLL by 0.02–0.97; their most frequent token frequency was at most 1.76/D, normalized marginal entropy at least 0.94, and at least 98% of stored sequences were unique.The uniform baseline is ln D, the NLL of a model assigning probability 1/D to every token; accepted teachers had consecutive-token mutual information 0.008–0.90 and teacher perplexity 1.5–15.7.Source: §G.1 · pp. 83–84
  • Federated language-modeling experiments use four SlimPajama domains to instantiate four clients, excluding the inaccessible book domain; allocations are IID (β=0.0), mixed (β=0.7), or pure non-IID (β=1.0).Source: §G.1 · pp. 84–85

G.2 MODELS

  • Federated training uses Mamba2 language models from scratch at depths 4 or 16, while model merging evaluates public 130M, 370M, and 780M Mamba2 checkpoints.The depth-4 backbone ablation replaces some or all Mamba2 mixers with causal attention layers; the attention-only model is the zero-SSM control. Synthetic learners match their teacher’s single-layer architecture and dimensions but do not receive teacher parameters.Source: §G.2 · pp. 86–87

G.3 SYNTHETIC DISTRIBUTED LEARNING FRAMEWORK

  • The synthetic framework runs centralized learner training, gradient-bound and smoothness evaluations at checkpoints, and FedAvg and FedProx training for each configuration and learner seed.Centralized training minimizes teacher-forced next-token cross-entropy and projects parameters back into the stability set after every update; evaluations restart from the zero recurrent state for each stored sequence.Source: §G.3 · p. 88
  • Synthetic evaluations check gradient and smoothness bounds on fixed-seed test sequences and compare FedAvg and FedProx under IID and dominant-token client partitions.Federated runs use equal-size shards, with each client performing E local SGD updates per round; the base setting covers 72 configurations and one-factor variations cover 16 core configurations.Source: §G.3 · p. 89

4 GSSM(W r)2 P

  • The GSSM evaluation sets reference radii and envelopes from initialization, then compares theorem bounds with measured subset-gradient estimates while logging iterate containment and projection activity.For FedProx, it uses a reduced radius and corresponding envelopes.Source: §4 · p. 90

G.4 PRACTICAL DISTRIBUTED LEARNING FRAMEWORK

  • The practical framework compares nine federated algorithms under a shared training schedule, with hyperparameters selected by validation perplexity averaged over the final 10 rounds and three seeds.It also reports train and validation perplexities and includes client-count and backbone ablations.Source: §G.4 · pp. 90–92

G.5 MODEL MERGING FRAMEWORK

  • Model merging fine-tunes a pretrained base independently on four downstream tasks, then merges the resulting specialists and evaluates each merged model on every task’s evaluation split.Fine-tuning masks the causal-language-model loss to label-token positions; the evaluation-pool subsample is fixed across seeds for comparability.Source: §G.5 · pp. 92–93

G.6 REPRODUCIBILITY

  • Experiments use single-GPU A100-40GB and H100 workers, with main federated and merging sweeps run on H100s and tokenized data cached for reuse.Federated runtime scales approximately linearly with local steps E, ranging from about 30 minutes at E = 64 to roughly 2 hours for E = 256 runs on an H100.Source: §G.6 · p. 93

H.1 SINGLE LAYER SYNTHETIC DISTRIBUTED LEARNING

The single-layer synthetic experiments find no gradient-bound violations, while showing that recurrent memory and federated settings shape bound behavior and optimization quantities. The numerical validation is limited to single-layer theory; the supplemental practical experiments cover nine algorithms across 72 configurations.

H.1 SINGLE LAYER SYNTHETIC DISTRIBUTED LEARNING

  • The multi-layer bound composes per-layer constants with layerwise sensitivities, but numerical evaluation validates only the single-layer theory.The composition does not establish greater relative looseness; useful layers may also contract perturbations.Source: §H.1 · pp. 93–94
  • No violations occurred, and per-sequence tightened-bound utilization remained at least two orders of magnitude below the bound at every ρ.The evaluation reports zero violations for configured and tightened bounds and shows utilization distributions ending at least two orders below the bound.Source: §H.1 · p. 94, Table 17 · p. 94, Table 18 · p. 94
  • Certified dimension exponents follow the N D^2-type structure, whereas the measured gradient is flat in N and decreases with D.The reported exponents are from one-factor fits at the final learner checkpoint.Source: §H.1 · p. 94, Table 20 · p. 96
  • Learners generally outperform uniform prediction and approach teacher test NLL; at long memory they nearly match teachers and train close to the stability limit.At short memory, the teacher’s small predictive advantage over uniform makes normalized gap recovery sensitive to small absolute errors.Source: §H.1 · p. 95, Table 21 · p. 96, Figure 5 · p. 97
  • Client drift grew as E^0.87 versus the reference radius’s E^1, while disagreement grew as K^1.48 and (1−ρ)^−1.15, faster than squared gradient norm’s (1−ρ)^−0.2.Measured disagreement was 10^-11 to 10^-14 of the certified bound; the sequence-level squared-TV aggregate was unchanged by the dominant-token partition, unlike the token-level analogue.Source: §H.1 · p. 95, Figure 6 · p. 97
  • The recurrent-history term carried at least 99% of the bound; its memory factor contributed 93% of log-bound variation, while the transition factor offset it by 13%.For utilization variation, memory contributed 102% and measured gradient growth −18%; recurrence dominated variation in tightness, not necessarily absolute looseness.Source: §H.1 · p. 96, Table 22 · p. 96

H.2 PRACTICAL DISTRIBUTED LEARNING

  • The practical supplement evaluates nine federated algorithms across 72 main configurations and adds client-count and backbone ablations.The main configurations vary local rounds, model depths, heterogeneity, and data composition.Source: §H.2 · pp. 96–97

H.2.1 DETAILED RESULTS ACROSS THE 72 MAIN CONFIGURATIONS

Across the 72 main configurations, FedProx ranks first in nearly all reported metric columns, while heterogeneity and local-step count alter the relative gains of other methods. Late-training rankings contrast with FedDyn’s and FedAdam’s all-round validation performance.

H.2.1 DETAILED RESULTS ACROSS THE 72 MAIN CONFIGURATIONS

  • FedProx has the lowest mean perplexity in 34 of 36 metric columns; the exceptions are β = 1 last-10 validation, where FedDyn is 40.0 versus 40.2, and E = 64 all-round validation, where SCAFFOLD is 105.6 versus 106.6.On last-10 perplexity, FedProx, FedDyn, and SCAFFOLD rank first, second, and third across local-step counts, domain settings, and model depths.Source: §H.2.1 · p. 97, Table 23 · p. 99
  • FedProx’s last-10 perplexity is approximately 19%, 20%, and 34% below FedAvg at β = 0, 0.7, and 1; SCAFFOLD rises from eighth to third place as β increases.At β = 1, SCAFFOLD improves on FedAvg by 18%, while FedExP has 13.5% higher last-10 perplexity than FedAvg.Source: §H.2.1 · pp. 97–98, Table 23 · p. 99
  • As E rises from 64 to 256, FedDyn’s and SCAFFOLD’s last-10 reductions versus FedAvg grow from about 12% and 4.5% to 22% and 10%; FedProx stays at 24–27%.At fixed 50 communication rounds, all nine algorithms achieve approximately 35–45% lower last-10 perplexity at E = 256 than at E = 64.Source: §H.2.1 · p. 98
  • FedAlign and FedSAM stay within 2% of FedAvg and DiLoCo-style within 5% on last-10 perplexity, while FedAdam is approximately 9% higher across local-step counts and depths.FedAdam’s gaps across domain settings are approximately 5–12%.Source: §H.2.1 · p. 98
  • Despite strong late-training results, FedDyn’s and FedAdam’s all-round validation averages are about 1.7–6.4 times FedAvg’s; FedDyn finishes 18% below FedAvg, but FedAdam finishes 9% above.After round 1, their validation perplexities are about 8 and 9 times FedAvg’s, respectively; FedDyn falls below FedAvg from round 4 onward.Source: §H.2.1 · p. 99
  • Domain setting changes absolute last-10 perplexity without changing the three leading methods’ ordering, while greater depth narrows their relative gains over FedAvg.FedProx’s improvement over FedAvg narrows from 28% to 22% with depth, FedDyn’s from 22% to 12%, and SCAFFOLD’s from 9% to 4%.Source: §H.2.1 · p. 99

H.2.2 CLIENT-COUNT ABLATION

In the evaluated client-count protocol, the three strongest methods from the main grid remain strongest, but their relative gains and the ordering of other methods vary with K. The ablation changes several aspects of the training setup, so it does not isolate client count as a cause.

H.2.2 CLIENT-COUNT ABLATION

  • The ablation varies K across 4, 8, 16, and 32 at β = 0.7, depth 4, and E = 64, averaging over four domain settings.Each client holds a 4/K share of the same domain pools.Source: §H.2.2 · p. 99
  • With per-client work held fixed, FedAvg’s last-10 validation perplexity decreases by approximately 4% from K = 4 to K = 8, from 64.2 to 61.6, with little further change for larger K.Source: §H.2.2 · p. 99
  • FedProx, FedDyn, and SCAFFOLD remain the strongest on last-10 validation perplexity at every K; FedProx reduces perplexity versus FedAvg by 22–24% across evaluated client counts.SCAFFOLD’s reduction grows from 3% to 13%, and FedDyn’s from 7% to 18%, between K = 4 and K = 32.Source: §H.2.2 · p. 100
  • The results describe the evaluated protocol at one β, depth, and E; changes in total client work, local dataset sizes, and partitioning prevent isolating client count’s effect.The proposed stable-aggregate-update explanation is not directly tested.Source: §H.2.2 · p. 100

H.2.3 MODEL BACKBONE ABLATION

The backbone ablation finds that both attention-layer placement and federated algorithm choice affect validation perplexity. In one-shot Mamba2 merging, Task Arithmetic and TIES-Merging outperform weight averaging, while the analysis highlights a trade-off between neighborhood containment and useful adaptation.

H.2.3 MODEL BACKBONE ABLATION

  • FedProx is best on all six backbones at 17–21% lower perplexity than FedAvg, while FedAlign and FedSAM remain within 3% of FedAvg.SCAFFOLD gains 13–15% on hybrids with two or three Mamba2 layers but only 4% on Mamba2 and 1% on attention-only; FedDyn is unstable across architectures.Source: §H.2.3 · p. 101
  • As Mamba2 layers are replaced, server-step penalties grow: FedAdam rises from +8% on Mamba2 to +40% on attention-only, while DiLoCo-style rises from +1% to +14%.FedExP is neutral on Mamba2 but 3–6% better than FedAvg on every backbone containing attention.Source: §H.2.3 · p. 101

H.3 MODEL MERGING

  • After four epochs, specialists reach 89–91% accuracy on their own tasks but transfer little to the other tasks; larger checkpoints narrow, but do not eliminate, deficits to specialists.The reported remaining deficits are especially pronounced on DBpedia and TREC.Source: §H.3 · pp. 101–102, Table 30 · p. 103

H.2.3 MODEL BACKBONE ABLATION

  • 26.5 last-10 validation perplexity for hybrid75 ranks best under FedAvg, versus 30.7 for Mamba2; placing one attention layer in the middle improves on pure Mamba2 by 14%.Attention placed first or a mostly-attention stack loses that gain; attention-only is 38% worse than pure Mamba2 at this model size and sequence length.Source: §H.2.3 · p. 100, Table 28 · p. 102

H.3 MODEL MERGING

  • 61.3%, 71.8%, and 72.1% task-averaged accuracy for the 130M, 370M, and 780M checkpoints beat weight averaging at 37.4%, 50.4%, and 51.7%, respectively.The best evaluated merges use Task Arithmetic or TIES-Merging.Source: §H.3 · pp. 101–102, Table 29 · p. 102
  • Task Arithmetic sweeps favor λ = 0.7–1 over the averaging scale λ = 0.25, while the best evaluated four-epoch TIES setting is (λ, d) = (1.3, 0.1).For Mamba2-130M, the proximal penalty reduces four-epoch averaging accuracy from 37.4% at µ = 0 to 26.2% at µ = 10^-1.Source: §H.3 · pp. 102–103, Table 32 · p. 103, Table 33 · p. 103
  • Task-vector alignment is low (∥P_F∥ = 1.04–1.07 for the 780M model); the authors suggest averaging can suppress useful adaptation, and uniform smoothness does not explain larger DBpedia and TREC gaps.One-shot merging has no later rounds to recover adaptation suppressed by limiting specialist drift.Source: §H.3 · p. 102

Related papers

Loading 2610.02659v1…