Paper deep dive

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

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

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

Intelligence

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

Last extracted: 10/7/2026, 7:57:52 AM

Summary

This paper presents an architecture-aware convergence analysis for distributed learning with selective State Space Models (SSMs), specifically Mamba2. It derives gradient and smoothness bounds that account for SSM-specific properties like recurrent stability and input-dependent discretization, and establishes convergence bounds for FedAvg and FedProx. The theoretical findings are validated numerically and through experiments comparing nine federated learning algorithms across six text domains.

Entities (9)

Adam Piaseczny · person · 99%FedAvg · algorithm · 99%FedProx · algorithm · 99%Purdue University · organization · 99%Mamba2 · model-architecture · 98%Selective State Space Models · model-class · 97%Federated Learning · task · 95%Mamba · model-architecture · 95%S4 · model-architecture · 90%

Relation Signals (6)

Adam Piaseczny → affiliatedwith → Purdue University

confidence 99% · Adam Piaseczny Affiliation: Purdue University

FedAvg → isusedfor → Federated Learning

confidence 99% · convergence bounds for FedAvg... federated learning algorithms

FedProx → isusedfor → Federated Learning

confidence 99% · convergence bounds for FedAvg and FedProx

Mamba2 → isa → Selective State Space Models

confidence 98% · Modern state space models (SSMs), such as Mamba2... characterize modern selective SSMs.

Selective State Space Models → hascomponent → Mamba

confidence 85% · modern SSMs, such as Mamba... have emerged as one of the most promising new paradigms.

Selective State Space Models → hascomponent → S4

confidence 85% · modern SSMs... include... Structured State Space (S4) sequence model

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

418,916 characters extracted from source content.

Expand or collapse full text

Distributed Learning with Selective State Space Models: Architecture-Aware Convergence Analysis Adam Piaseczny Affiliation: Purdue University Email: apiasecz@purdue.edu Md Kamran Chowdhury Shisher Affiliation: University of Colorado Colorado Springs Email: mshisher@uccs.edu Shiqiang Wang Affiliation: University of Exeter Email: shiqiang.wang@ieee.org Christopher G. Brinton Affiliation: Purdue University Email: cgb@purdue.edu 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. 1 Introduction The modern generative AI landscape is largely dominated by transformer architecture (Vaswani et al., 2017), with the majority of the state-of-the-art large language models (LLMs) employing them (Agarwal et al., 2025; Team et al., 2026a; Qiu et al., 2026; Team et al., 2026b). Despite this widespread success, transformers face critical limitations in certain scenarios, often driven by the core efficiency bottleneck of their self-attention mechanism, scaling quadratically as ⁡(L2)O(L^2) with sequence length L in both time and memory complexity (Tay et al., 2022). To address this inefficiency, a variety of alternative mechanisms and architectures have been proposed, including sparse attention (Child et al., 2019; Zaheer et al., 2020), linear attention approximations (Katharopoulos et al., 2020; Choromanski et al., 2021), and modernized recurrent frameworks (Peng et al., 2023). Among these innovations, state space models (SSMs) (Gu et al., 2022b) have emerged as one of the most promising new paradigms. These modern SSMs, such as Mamba (Gu and Dao, 2024), overcome the quadratic bottleneck by operating with a strictly linear-time complexity, ⁡(L)O(L), during training, and a constant-sized latent state, ⁡(1)O(1). Driven by these computational advantages, the SSM architecture has undergone rapid development, achieving competitive performance at a reduced computational cost compared to transformers. Architectural innovations, such as the Structured State Space (S4) sequence model, HiPPO initialization (Gu et al., 2022b), and selective mechanisms (Gu and Dao, 2024) have enabled SSMs to dynamically filter information and capture and retain long-range dependencies. Most recently, SSMs have demonstrated capacity for achieving state-of-the-art language modeling performance at scale (Lahoti et al., 2026; Zuo et al., 2024; Blakeman et al., 2025; Team et al., 2024). Research Gap. Despite these advancements, the deployment of SSMs in distributed learning environments, where training is conducted across a set of agents, remains largely unexplored. While some studies have considered the intersection of SSMs with federated learning (FL) (McMahan et al., 2017), most of these frameworks are architecture-agnostic. Conversely, existing SSM-specific distributed learning studies either focus on classical, non-selective SSM settings (Keçeci et al., 2025), or use SSMs for specialized personalized-FL applications (Zhang et al., 2026), leaving the convergence behavior of modern selective SSMs largely uncharacterized. Developing a foundational understanding of selective SSM convergence behavior in distributed settings can enable architecture-aware parameter tuning, protocol selection, and other control tasks for training performance optimization, uncovering another advantage of SSMs over transformer LLM approaches that rely on architecture-agnostic optimization tools (Fang et al., 2026; Soro et al., 2026). We thus pose the following question: How can we leverage the structural properties of modern selective SSMs to derive architecture-aware convergence bounds that rigorously characterize training behavior in terms of core state space parameters for distributed and federated learning techniques? Contributions. In answering this question, our main contributions are summarized as follows: • Convergence analysis for distributed selective SSMs: We provide an architecture-aware gradient norm bound for selective SSMs, revealing the impact of SSM parameters like core state-space parameter blocks. Then, we use this bound to establish the convergence rate of popular FL algorithms training selective SSMs, showing dependence on standard distributed learning factors like smoothness as well as architecture-specific parameters like block-wise spectral-norm envelopes. To the best of our knowledge, our work is the first to provide an architecture-aware convergence rate of FL algorithms for selective SSMs. • SSM theory validation and theory-motivated algorithm comparison: We validate our theoretical results using teacher-model derived sequences and single-layer selective SSMs, observing no gradient-bound violations, increased client disagreement as ρ approaches one, and increased drift with more local steps. We use the theoretical takeaways to rank federated algorithms on heterogeneous-domain datasets. Our results show that these rankings highly correlate with validation performance, supporting the practical relevance of the SSM-specific quantities identified by our analysis. 2 Related Work State Space Models (SSMs). SSMs originate from classical control theory concepts, where they model dynamic systems through the evolution of latent states over time (Kalman, 1960). The modern SSMs were introduced with the pivotal Structured State Space (S4) (Gu et al., 2022b) and its variants (Gu et al., 2022a; Hasani et al., 2023; Gu et al., 2020), and leverage structured parameterizations to efficiently capture long-range dependencies. Selective SSMs introduced input-dependent state-space parameters to dynamically filter information (Gu and Dao, 2024) while the Structured State Space Duality (SSD) established connections with linear attention and enabled efficient implementations (Dao and Gu, 2024), enabling better expressivity, training stability, and hardware-aware optimizations (Hwang et al., 2024; Kim et al., 2026; Lahoti et al., 2026). These advances have driven growing adoption across language modeling (Waleffe et al., 2024; Blakeman et al., 2025; He et al., 2025; Glorioso et al., 2024; Ren et al., 2025), computer vision (Li et al., 2024; Zhu et al., 2024; Xie et al., 2025), and other sequence modeling applications (Ma et al., 2025). Distributed Learning with SSMs. Distributed learning enables model training across computational nodes, with federated learning (FL) addressing settings where data remain distributed among clients. Since the introduction of a foundational federated learning algorithm - FedAvg (McMahan et al., 2017) - numerous methods have been proposed to address statistical heterogeneity, communication costs, and client drift (Li et al., 2020; Karimireddy et al., 2020; Wang et al., 2020; Reddi et al., 2021; Acar et al., 2021; Li et al., 2021). However, despite the growing adoption of SSMs, theoretical understanding of their behavior in these settings remains limited. The existing SSM-related approaches address specific challenges: FedAlign (Keçeci et al., 2025) tackles the problem of models converging to different parameter basins in federated system identification, while FedSSM (Zhang et al., 2026) uses SSMs as a meta-learning tool within personalized FL. Related works also explore merging operations on hidden states, contexts, or representations (Liu et al., 2025; Pióro et al., 2024), rather than weight-space aggregation of independently trained selective SSMs. As a result, there remains an open question of how selective SSM dynamics affect optimization and aggregation guarantees in distributed learning settings. This motivates our study of connecting the architectural properties of selective SSMs to their behavior in distributed training. 3 System Model (a) Selective SSM layer (b) Distributed SSM training 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 E local updates on local datasets kD_k, and the resulting models are aggregated into Wr+1W^r+1. State Space Model. We consider a decentralized learning environment operating over a global selective state space model (SSM) parameterized by W=(WB,WC,WΔ,A)W= (W_B,W_C,W_ ,A ) where WΔ∈ℝD×DW_ ^D× D, WB∈ℝN×DW_B ^N× D, WC∈ℝN×DW_C ^N× D, and A∈ℝN×1A ^N× 1, with A representing diagonal entries of a structured matrix =diag(A)A=diag(A). The SSM architecture is inspired by classical control theory concepts, in particular modeling dynamic systems through the evolution of latent states over time. These systems map a continuous input signal x⁡(t)∈ℝx(t) to an output sequence y⁡(t)∈ℝy(t) through an evolving, compressed latent state h⁡(t)∈ℝNh(t) ^N using linear time-invariant (LTI) differential equations: However, since neural networks process discrete sequences, these continuous-time SSM dynamics are mapped to a discrete recurrence using a discretization rule. Under this discrete formulation, the model maps an input sequence x=(x1,…,xS)x=(x_1,…,x_S) where xt∈ℝD×1x_t ^D× 1 to a sequence of outputs yt∈ℝD×1y_t ^D× 1 through a sequence of compressed latent states ht∈ℝN×Dh_t ^N× D. While earlier foundational models utilized various discretization rules (Gu et al., 2022b), we adopt the first-order exponential-Euler discretization used in the more recent architectures (Lahoti et al., 2026). This discretization scheme retains the exponential state transition while applying a first-order approximation to the input contribution. We analyze the discrete model defined by this recurrence. In addition to this discrete foundation, modern SSMs introduce a selective mechanism that enables dynamic filtering of relevant information across the input sequence context. The selective SSMs achieve this by making the state-space and time-scale parameters depend on the current input. Specifically, given the model weights W, the input-dependent SSM quantities corresponding to an input token xtx_t at the sequence position t are defined as: • Time-scale parameter: Δt=Softplus(WΔxt) _t=Softplus(W_ x_t), where Δt∈ℝD×1 _t ^D× 1, acts as a gating mechanism, controlling the discretization step size to determine importance of the current token. • State transition matrix: A¯t=exp⁡(AΔt⊤) A_t= (A _t ), where A¯t∈ℝN×D A_t ^N× D, determines the retention and decay of the historical context, managing the system’s memory across the sequence. • Input matrix: B¯t:=BtΔt⊤=(WBxt)(Softplus⁡(WΔxt))⊤ B_t:=B_t _t =(W_Bx_t)(Softplus(W_ x_t)) , where B¯t∈ℝN×D B_t ^N× D, and Bt:=WBxtB_t:=W_Bx_t. Under the adopted exponential-Euler discretization, this matrix modulates how strongly the current input vector influences the newly updated latent state. • Output matrix: Ct=WCxtC_t=W_Cx_t, where Ct∈ℝN×1C_t ^N× 1, linearly projects the compressed latent state back into the target output space. Given these dynamically generated matrices, the discrete hidden state sequence hth_t and the corresponding output yty_t at each sequence step are computed according to the following recurrence: where h0=N×Dh_0=0_N× D, the exponential and Softplus functions are applied element-wise, and ⊙ denotes the Hadamard product. With the architecture specified above, we next define the local optimization objective. Let ξ=(x,y∗)ξ=(x,y^*) denote an input-target sequence pair, where x=(x1,…,xS)x=(x_1,…,x_S) is the input sequence and y∗=(y1∗,…,yS∗)y^*=(y_1^*,…,y_S^*) is the corresponding target sequence. The sequence loss is then the average token loss, defined as: ℒS(W,ξ):=1S∑t=1Sℓt(yt(W,x),yt∗) ^S (W;ξ ):= 1S _t=1^S _t (y_t(W;x),y_t^* ) (3) where yt(W,x)y_t(W;x) is the model prediction at position t parameterized by W, and ℓt _t is some differentiable loss function. When the sample ξ is fixed, we write ℒS(W)L^S(W) and ℓt(yt) _t(y_t) for notational simplicity. Figure 1(a) illustrates the resulting architecture, the parameter transformations, and state recurrence. Distributed Optimization Environment. Given the selective SSM setup established above, we consider the deployment of this architecture in a decentralized learning setting. In this framework, a central server initializes a global selective SSM model parameterized by weights W0W^0, which is collaboratively trained across a network of K clients. Each client k∈1,…,Kk∈\1,…,K\ possesses a local dataset k=ξk,i∼ki=1|k|D_k= \ _k,i _k \_i=1^|D_k|, where kP_k denotes the local data distribution of client k and each sample ξk,i=(xk,i,yk,i∗) _k,i=(x_k,i,y_k,i^*) represents a discrete input sequence and its corresponding target sequence. For autoregressive next-token prediction, the target at each sequence position satisfies yt∗=xt+1y^*_t=x_t+1. Given this setup, the framework aims to construct a global model that incorporates the combined knowledge of these diverse local distributions. Formally, the global objective is to find the optimal global model parameters W∗W^* that minimize the aggregated expected sequence modeling loss: W∗=arg⁡minW⁡f⁡(W),wheref⁡(W)=∑k=1Kpkℒk(W) W^*= _Wf(W), f(W)= _k=1^Kp_kL_k(W) (4) where pkp_k are positive weights satisfying ∑k=1Kpk=1 _k=1^Kp_k=1, and ℒk(W)=ξ∼k[ℒS(W,ξ)]L_k(W)=E_ξ _k [L^S(W;ξ) ] is the expected sequence loss for client k over its local distribution. To solve this optimization problem, the framework employs a decentralized training procedure initialized with the global model W0W^0. The optimization then proceeds over discrete communication rounds r∈:=0,1,…,T−1r := \0,1,…,T-1 \. At each round r, the server shares the current global parameters WrW^r with the participating clients. Each client k initializes its local model as Wkr,0=WrW_k^r,0=W^r and performs E local steps of stochastic gradient descent (SGD) with step size η to minimize a local objective Fk(W)F_k(W). The specific formulation of Fk(W)F_k(W) depends on the distributed learning algorithm employed. For instance, in the classical FedAvg algorithm (McMahan et al., 2017), the local objective of client k is the expected sequence loss, i.e., Fk(W):=ℒk(W)F_k(W):=L_k(W). In contrast, FedProx (Li et al., 2020) modifies the local objective as Fkr(W):=ℒk(W)+μ2‖W−Wr‖F2F_k^r(W):=L_k(W)+ μ2 \|W-W^r \|^2_F, with μ>0μ>0, to mitigate local divergence caused by data heterogeneity. Through this local optimization step, the features of different distributions kP_k are implicitly embedded into the updated local weights, resulting in a set of specialized models Wkr,EW_k^r,E. Following the completion of the E local training steps, the server combines these localized models into a unified global representation using a generalized aggregation function: Wr+1=Aggregate(Wkr,Ek=1K,Wr)W^r+1=Aggregate ( \W_k^r,E \_k=1^K;W^r ). In standard distributed baselines such as FedAvg and FedProx, this aggregation is implemented as a simple data-weighted linear average: Wr+1=∑k=1K|k|∑j=1K|j|Wkr,E W^r+1= _k=1^K |D_k| _j=1^K|D_j|W_k^r,E. By repeating this process over T communication rounds, the framework iteratively refines the global model toward an optimal parameter set W∗W^* by aggregating information learned across the participating clients. Figure 1(b) illustrates this iterative training process. 4 Theoretical Analysis Given this decentralized SSM setup, we consider distributed training procedures that update a global model through repeated local adaptation and server-side aggregation. These procedures induce a global optimization trajectory Wrr=0T\W^r\_r=0^T over the objective defined in (4). Since f is generally nonconvex, we characterize convergence through the standard expected gradient stationarity criterion 1T∑r=0T−1‖∇f(Wr)‖F2. 1T _r=0^T-1E \|∇ f(W^r) \|_F^2. The central problem is to derive convergence bounds that explicitly capture the structure of selective SSMs, rather than relying solely on architecture-agnostic smoothness, variance, and heterogeneity constants. In particular, the recurrent state update induce dependencies between the local optimization trajectory and SSM-specific quantities including the stability of the discretized transition matrices, the input-dependent time scale parameter Δt=Softplus⁡(WΔxt) _t=Softplus(W_ x_t), and the norms of the state-space projection parameters. Our convergence analysis characterizes how these SSM-specific quantities control gradient growth, client drift, and aggregation error. As such, in the remainder of this section, we establish theoretical guarantees for federated training of selective SSMs. We first derive an explicit gradient norm bound for a single-layer selective SSM and use it to characterize federated optimization in terms of SSM-specific parameters. We provide convergence guarantees for FedAvg and FedProx, and relate additional federated baselines to the resulting stability, parameter-sensitivity, and distribution-shift quantities. The main text highlights the key results and their implications, while full proofs, supplementary corollaries, and additional baseline analyses are provided in Appendices A, B, C, D, E. 4.1 Gradient Norm Bound for Selective SSMs The main technical challenge lies in controlling the parameter sensitivity induced by the input-dependent SSM recurrence. The loss gradient depends on the model parameters both directly through the output projection and indirectly through their influence on the recurrent hidden states. By bounding these two contributions separately and subsequently combining them, we derive a sequence-wise gradient bound that is independent of the sequence length. We first introduce the notation used in our analysis. An input-target pair ξ is admissible if ξ∈ξ , where Z is a fixed, parameter-independent measurable set of input-target sequences satisfying Prξ∼k⁡(ξ∈)=1Pr_ξ _k (ξ )=1 for all clients k. All conditions below hold uniformly for admissible sequences, finite lengths S≥1S≥ 1, and token indices t≤St≤ S, with finite constants independent of S, ξ, and t. Next, we state the necessary definitions and assumptions on the loss and the input sequence. Definition 4.1 (SSM Stability). For γ,δ>0γ,δ>0, define γ,δS_γ,δ as the set of parameter tuples W satisfying An≤−γA_n≤-γ for every n∈1,…,Nn∈\1,…,N\, where AnA_n is the n-th element of the transition parameter A, and [Softplus⁡(WΔxt)]d≥δ [Softplus(W_ x_t) ]_d≥δ for every feature coordinate d∈1,…,Dd∈\1,…,D\ and admissible input token. Consequently, every entry of A¯t A_t satisfies: (A¯t)n,d≤ρ:=exp⁡(−γδ)<1 ( A_t )_n,d≤ρ:= (-γδ)<1. Assumption 4.2 (Output loss regularity). The token loss ℓt(⋅,yt∗) _t(·,y_t^*) is twice continuously differentiable with respect to its prediction argument. For the parameter tuple W under consideration, there exist finite constants Lℓ,Hℓ>0L_ ,H_ >0 such that, for every admissible ξ, the prediction y=yt(W,x)y=y_t(W;x) satisfies ‖∇yℓt(y,yt∗)‖2≤Lℓ,‖∇y2ℓt(y,yt∗)‖op≤Hℓ. \| _y _t(y,y_t^*)\|_2≤ L_ , \| _y^2 _t(y,y_t^*)\|_op≤ H_ . (5) Note that for softmax cross-entropy over logits, the Hessian bound holds globally. Assumption 4.3 (Bounded Inputs). There exists Lx>0L_x>0 such that ∀t∀ t, ‖xt‖2≤Lx\|x_t\|_2≤ L_x ∀ admissible ξ. Theorem 4.4 (SSM Gradient Bound). Suppose Assumptions 4.2 and 4.3 hold for a parameter tuple W∈γ,δW _γ,δ. For a single-layer SSM defined in Section 3, the following bound holds for every finite integer S≥1S≥ 1 and every admissible length-S sequence pair ξ∈ξ : ‖∇WℒS(W,ξ)‖F≤SSM(W):=LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ) \| _WL^S(W;ξ) \|_F _SSM(W):=L_ L_x (K_BK_ ND1-ρ+\|W_C\|_2 K_W1-ρ ) (6) with KΔ K_ :=‖WΔ‖2Lx+Dlog⁡(2),KB:=‖WB‖2Lx2, := \|W_ \|_2L_x+ D (2),~~~~K_B:= \|W_B \|_2L_x^2, (7) KW K_W :=NDKBKΔ(KΔ+D‖A‖2Lx)(ρ1−ρ)+Lx2NKΔ+LxDKB :=NDK_BK_ (K_ + D \|A \|_2L_x ) ( ρ1-ρ )+L_x^2 NK_ +L_x DK_B (8) where ∥⋅∥F\|·\|_F is the Euclidean norm of the block Frobenius norms. Discussion. The key technical step is to bound both the hidden states and their parameter Jacobians under input-dependent transitions. The stability condition ρ=exp⁡(−γδ)<1ρ= (-γδ)<1 allows us to reduce the recurrent contributions to geometric sums, yielding an explicit parameter-gradient bound rather than requiring such a bound as an assumption. The resulting architecture-aware bound is independent of the sequence length while retaining explicit dependence on ρ and the core selective SSM parameter blocks WΔW_ , WBW_B, WCW_C, and A. The complete proof is provided in Appendix A. Appendix C extends the analysis to the multi-layer setting by combining stability with layerwise input-output sensitivities, resulting in the bound SSM(Q)(W)G_SSM^(Q)(W) stated in Theorem C.3. Together, these results provide architecture-specific control of the gradient norm required for the subsequent federated convergence analysis. 4.2 Convergence Analysis To characterize federated optimization with selective SSMs, we introduce a local parameter neighborhood and impose the stochastic-gradient conditions required for the convergence analysis. Definition 4.5 (R-Local Neighborhood). For a reference model W~ W and radius R>0R>0, define R(W~):=W:‖W−W~‖F≤RN_R( W):=\W:\|W- W\|_F≤ R\, where the ∥⋅∥F\|·\|_F is defined for a parameter tuple W as ‖W‖F:=(‖WΔ‖F2+‖WB‖F2+‖WC‖F2+‖A‖F2)1/2\|W\|_F:= (\|W_ \|_F^2+\|W_B\|_F^2+\|W_C\|_F^2+\|A\|_F^2 )^1/2. In the following convergence analysis, the regularity conditions of Assumption 4.2 are assumed to hold uniformly over each neighborhood R(Wr)N_R(W^r). In particular, the same finite constants LℓL_ and HℓH_ apply to every W∈R(Wr)W _R(W^r), every admissible input-target pair, and every token position. The stability margins γ,δγ,δ and input bound LxL_x are likewise fixed throughout the neighborhood. Assumption 4.6 (Unbiased stochastic gradients). For every client k and every stable local model neighborhood R(Wr)⊆γ,δN_R(W^r) _γ,δ, define gk(W,ξ):=∇WℒS(W,ξ)g_k(W;ξ):= _WL^S(W;ξ) (9) as the stochastic sequence-level gradient for a sampled sequence ξ∼kξ _k. Then, ∀W∈R(Wr)∀~W _R(W^r) ξ∼k[gk(W,ξ)]=∇Wℒk(W)E_ξ _k [g_k(W;ξ) ]= _WL_k(W) (10) holds. At each local update, conditional on all preceding training randomness, the mini-batch consists of Ξ independent sequences drawn from kP_k. The mini-batch gradient is their average. Batches drawn by different clients at the same local step are conditionally independent. Lemma 4.7 (Local SSM smoothness). Let WrW^r be a fixed reference parameter tuple and let 0<R<∞0<R<∞ satisfy R(Wr)⊆γ,δN_R(W^r) _γ,δ. Suppose Assumptions 4.2, 4.3, and 4.6 hold, with finite and differentiable expected local objectives ℒkL_k on this neighborhood. Furthermore, let the bounds in Assumption 4.2 hold with the same finite constants LℓL_ and HℓH_ ∀ W∈R(Wr)W _R(W^r), every admissible input-target sequence pair ξ∈ξ , and every token position, independently of sequence length. Then, for every client k, ℒkL_k is LSSM,RL_SSM,R-smooth on R(Wr)N_R(W^r). Specifically, for all W,V∈R(Wr)W,V _R(W^r), ‖∇Wℒk(W)−∇Wℒk(V)‖F≤LSSM,R‖W−V‖F,\| _WL_k(W)- _WL_k(V)\|_F≤ L_SSM,R\|W-V\|_F, (11) where LSSM,RL_SSM,R is defined in (129). With these conditions in place, we characterize the convergence of FedAvg for selective SSMs. The corresponding analysis for FedProx and supplementary corollaries are available in Appendix B, D. Corollary 4.8 (Single-layer FedAvg convergence rate). Under Assumptions 4.2- 4.6, the conditions specified in Definitions 4.1- 4.5, and the conditions of Theorem B.5, let MB,MC,MΔ,M_B,M_C,M_ , and MAM_A denote uniform upper bounds on ‖WB‖2,‖WC‖2,‖WΔ‖2, \|W_B \|_2, \|W_C \|_2, \|W_ \|_2, and ‖A‖2 \|A \|_2, respectively, over the R-local neighborhoods R(Wr)N_R(W^r), and define MR:=(1+MB)(1+MC)(1+MΔ)(1+MΔ+MA).M_R:=(1+M_B)(1+M_C)(1+M_ )(1+M_ +M_A). We further define ¯:=∑k=1Kpkk P:= _k=1^Kp_kP_k, ℋTV:=∑k=1KpkTV⁡(k,¯)2,H_TV:= _k=1^Kp_kTV(P_k, P)^2, where TVTV denotes the total variation distance. Under these conditions, for FedAvg applied to the single-layer selective SSM defined in Section 3, the global optimization trajectory Wrr=0T\W^r\_r=0^T satisfies 1T∑r=0T−1‖∇f(Wr)‖F2=⁡(f⁡(W0)−f⁡(W∗)ηET+ηMR4Ξ(1−ρ)8+η2MR6(E−1)(1−ρ)12(1Ξ+EℋTV)). 1T _r=0^T-1E \|∇ f(W^r) \|_F^2=O ( f (W^0 )-f (W^* )η ET+ η M_R^4 (1-ρ)^8+ η^2M_R^6(E-1)(1-ρ)^12 ( 1 +EH_TV ) ). The quantities N,D,Lx,Lℓ,N,D,L_x,L_ , and HℓH_ affect only the implicit constant and do not alter the displayed scaling with the optimization parameters, parameter envelopes, stability, or sequence length. Proof sketch and discussion. Standard FedAvg convergence bounds assume smoothness and bounds on stochastic gradient variance and client-gradient disagreement. In contrast, our analysis derives these quantities from the selective SSM recurrence. We first bound the hidden states and their parameter derivatives to obtain the gradient bound. The second-order analysis then establishes local smoothness, while the remaining lemmas translate the gradient bound into bounds on stochastic variance and client disagreement. Together, these results yield a convergence rate whose model dependence is explicit in the parameter envelopes MB,MC,MΔ,MAM_B,M_C,M_ ,M_A and the stability factor ρ, with dimensions and input and loss regularity constants fixed. As a result, the corollary identifies which SSM parameters enter the convergence guarantee without requiring a separately assumed smoothness or gradient-dissimilarity constant for the model. The derived bound distinguishes sequence length from recurrent memory: the gradient and smoothness estimates contain no explicit sequence-length factor, but increase as the recurrent stability factor ρ approaches one, allowing past state information to persist for longer. This dependence is stronger in the additional local-drift terms, which contain MR6/(1−ρ)12M_R^6/(1-ρ)^12, as opposed to the baseline stochastic term, which contains MR4/(1−ρ)8M_R^4/(1-ρ)^8. Consequently, the bound assigns a larger penalty to performing multiple local updates when ρ is close to one. However, ρ alone does not determine this penalty: even for the same ρ and client distributions, larger parameter envelopes increase the coefficient multiplying ℋTVH_TV. The analysis therefore separates the amount of distribution mismatch from the influence of the SSM parameters on its contribution to the convergence bound. This heterogeneity contribution vanishes at E=1E=1, but does not decrease with batch size when E>1E>1. Our synthetic experiments observe increasing client disagreement as ρ approaches one, supporting this qualitative connection without establishing the worst-case exponents. Furthermore, the radius R determines how the broadcast model’s weight norms are enlarged to obtain bounds valid throughout local training: each block satisfies ‖WX‖2≤‖WXr‖2+R\|W_X\|_2≤\|W_X^r\|_2+R. As such, allowing greater client displacement increases the parameter envelopes entering the convergence bound, even when ρ remains unchanged. Additional local steps can therefore affect the guarantee through both the explicit dependence on E and the larger envelopes needed to cover the resulting trajectories. These trajectories must also remain within a neighborhood satisfying R(Wr)⊆γ,δN_R(W^r) _γ,δ. The FedProx extension in Appendix B demonstrates how restricting client movement changes the guarantee: its effective radius Rμ≤R_μ≤ R permits smaller envelopes MXμ=MX−R+RμM_X^μ=M_X-R+R_μ, reducing the SSM-dependent factors multiplying noise and heterogeneity. However, the proximal term also increases local smoothness by μ. The resulting bound therefore exposes a trade-off between tighter control of the SSM parameters and the additional smoothness cost of proximal regularization. 5 Experiments In this section, we empirically evaluate the derived bounds and their implications for distributed SSM training in two parts: (i) direct evaluation of gradient and smoothness bounds using sequences generated by frozen SSM teachers, and (i) Mamba2 language modeling experiments with nine federated learning algorithms. In the latter, we compare theory-motivated expectations with observed performance illustrating the practical utility of the derived bounds. Details on our experimental setup are included in Appendix G, with supplementary experimental results available in Appendix H. Validation of Derived Bounds. We consider autoregressive token sequences generated by a frozen one-layer selective SSM parameterized by W=(WΔ,WB,WC,A)W= (W_ ,W_B,W_C,A ) analyzed in Theorem 4.4. This setup follows the formulation of Section 3 exactly, with the teacher SSM generating one-hot token (Lx=1L_x=1) sequences of length S=64S=64 for different combinations of D,ND,N and stability limits ρ, yielding 7272 distinct datasets. For each such dataset, we train a learner SSM on 40964096 sequences using projected SGD with softmax cross-entropy loss (Lℓ=2L_ = 2), with the model updates inside γ,δS_γ,δ. We check test-sequence gradients at five learner and one teacher checkpoints against the theorem’s SSMG_SSM. We also use SSMtightG^tight_SSM, tightened via realized stability, one-hot input bounds, and finite-history sums. We also evaluate the smoothness bound and examine the drift and disagreement quantities entering the FedAvg and FedProx analysis. Appendix H.1 provides additional validation results. 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). Statement Inequality checked Evaluations Violations Max. utilization Thm. 4.4 ‖∇ℒ‖F≤SSM\| \|_F _SSM, per sequence 1,327,104 0 / 0 1.0×10−31.0× 10^-3 / 8.3×10−38.3× 10^-3 Lemma 4.7 ‖∇2ℒ‖2^≤LSSM,0 \|∇^2L\|_2≤ L_SSM,0 41,472 0 7.0×10−57.0× 10^-5 Lemma 4.7 finite-difference ratio ≤LSSM,0.001≤ L_SSM,0.001 41,472 0 8.2×10−58.2× 10^-5 Lemma 4.7 ‖∇ℒ(Wa)−∇ℒ(Wb)‖F≤LSSM,R‖Wa−Wb‖F\| (W_a)- (W_b)\|_F≤ L_SSM,R\|W_a-W_b\|_F 864 0 1.0×10−51.0× 10^-5 Lemma B.1 |SSM(Wa)−SSM(Wb)|≤LSSM,R‖Wa−Wb‖F|G_SSM(W_a)-G_SSM(W_b)|≤ L_G_SSM,R\|W_a-W_b\|_F 864 0 0.420.42 (a) Gradient versus bound (b) Growth with memory (c) Grad. energy by WXW_X (d) Exponents in D and N 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 SSMG_SSM. (2(b)) At training end, SSMG_SSM and LSSM,RL_SSM,R scale as (1−ρ)−1.8(1-ρ)^-1.8 and (1−ρ)−3.6(1-ρ)^-3.6, versus (1−ρ)−0.5(1-ρ)^-0.5 and (1−ρ)−1.0(1-ρ)^-1.0 for measured gradient and estimated Hessian norms. (2(c)) As ρ→1ρ→ 1, gradient energy shifts from WB,WCW_B,W_C toward the transition block A. (2(d)) One-factor fits estimate the power-law exponent α for the derived (filled) and measured (open) gradient norms as vocabulary size D (blue) or state size N (orange) varies, with bootstrap 95% intervals. Results and Discussion. No violations were detected in 1,411,776 gradient, smoothness and bound-variation checks (Figure 2(a), Table 1), or in 103,200 federated checks across 1,200 FedAvg/FedProx runs on 72 teacher datasets. The federated runs explore step sizes beyond the theorems’ range. The results also reveal informative patterns. Memory amplifies gradient and curvature sensitivity. Across all generated datasets, the measured gradient and estimated Hessian norms grow approximately as (1−ρ)−0.5(1-ρ)^-0.5 and (1−ρ)−1.0(1-ρ)^-1.0, respectively, while their bounds grow as (1−ρ)−1.8(1-ρ)^-1.8 and (1−ρ)−3.6(1-ρ)^-3.6, as seen in Figure 2(b). Although the bounds grow more steeply, the curvature exponent is twice the gradient exponent for both the measured quantities and the certified constants. These trends are consistent with the recurrent-sensitivity structure of Theorem 4.4 and Lemma 4.7. The shared sensitivity factor Lx(HR+MCJR)L_x(H_R+M_CJ_R) enters the gradient bound linearly and the first term of the smoothness bound quadratically in (129). The approximately doubled fitted curvature exponent mirrors this linear-versus-quadratic dependence, although the observed growth is substantially milder than that of the bounds. Recurrent sensitivity and block-wise gradients. The recurrent-Jacobian term, which propagates parameter perturbations through the transition A¯t A_t, carries at least 99%99\% of SSMG_SSM at every ρ. The measured gradient energy (the share of ‖∇Wℒ‖F2\| _WL\|_F^2 carried by each block) shifts from WB,WCW_B,W_C (86%86\% at ρ=0.3ρ=0.3) to A (85%85\% at ρ=0.999ρ=0.999, Figure 2(c)). Dimension fits (Figure 2(d)) both vary with ρ and give bound exponents of 1.41.4-2.92.9 in D and 1.11.1-2.02.0 in N, reflecting the ND2N\,D^2-type factors in KWK_W. Appendix H.1 examines the quantities entering the convergence analysis. The observed E0.87E^0.87 growth of the ‖Wkr,E−Wr‖F\|W_k^r,E-W^r\|_F drift is consistent with the linear accumulation bound ηESSM,Rη EG_SSM,R used in the FedAvg analysis. Increasing client disagreement as ρ→1ρ→ 1 likewise agrees qualitatively with the SSM-sensitivity factor SSM,R2G_SSM,R^2 in Lemma B.4, which bounds disagreement by 4SSM,R2ℋTV4G_SSM,R^2H_TV. Across the bank, variation in log-utilization log10⁡(‖∇ℒ‖F/SSM) _10(\| \|_F/G_SSM) is dominated by the memory factor ρ/(1−ρ)2ρ/(1-ρ)^2. Its signed covariance shares are 93%93\% for log-bound variance and 102%102\% for log-utilization variance, the latter partly offset by measured gradient growth. Dimension and norm factors contribute less (Figure 9). This identifies the main source of variation in tightness, not absolute looseness. Limitations and outlook. The experiments did demonstrate some limitations of our derived bounds, mostly how conservative their values are in absolute terms, and how parameter-envelope growth with radius R over T rounds and saturation of the sequence-level total-variation metric affect the looseness. Nevertheless, the observations agree qualitatively with the analysis’s stability, sensitivity, and drift dependencies: no violations were detected, gradients and curvature grow with memory, blockwise gradients shift toward A, and drift grows with local steps. The recurrent factor ρ/(1−ρ)2ρ/(1-ρ)^2 dominates variation in tightness across the bank, with smaller contributions from projection and input factors. As such, they point to concrete future improvements: data-dependent constants that use the realized stability margin and the finite sequence length, envelopes that track the actual local trajectories rather than the worst case over the whole radius R, and a token-level dissimilarity that would make the heterogeneity term sensitive to the data partition. The Mamba2 experiments that follow examine whether the same drift and heterogeneity mechanisms help interpret practical federated training. Practical Federated Optimization. We train randomly initialized Mamba2 (Dao and Gu, 2024) models of depth 44 and 1616 across K=4K=4 clients for T=50T=50 rounds with E∈64,128,256E∈\64,128,256\ local steps and heterogeneity β∈0,0.7,1β∈\0,0.7,1\ on six RedPajama (Together Computer, 2023) domains under four domain assignments (β=0β=0 gives every client the same domain mix, β=1β=1 one domain each). Nine algorithms are compared over 7272 configurations and three seeds using mean last-ten-round validation perplexity, which also selects their hyperparameters (see Appendix G for details). Theory-Guided Expectations. An update rule affects radius R (in SSM,R,LSSM,RG_SSM,R,L_SSM,R), a heterogeneity term ∝η2LSSM,R2SSM,R2E(E−1)ℋTV η^2L_SSM,R^2G_SSM,R^2E(E-1)H_TV, and the server step to Wr+1W^r+1. For SGD, FedProx bounds drift by Rμ≤ηESSM,R_μ≤η EG_SSM,R (Theorem B.6 in Appendix B), at the cost of μ in Lprox,RμL_prox,R_μ, motivating benefits from tuned proximal regularization. SCAFFOLD and FedDyn can reduce corrected client disagreement when corrections track client-gradient differences, though benefits may be smaller when disagreement is small. FedExP, FedAdam and DiLoCo rescale server updates and may enlarge neighborhoods, while the norm-preserving FedAlign and small-perturbation FedSAM may behave like FedAvg, without a shared convergence guarantee (Appendix E). Table 2 summarizes these expectations for the AdamW-based Mamba2 experiments motivated by the insights provided by our analysis. Expected Δ (%) Method Theory-linked mechanism β=0β=0 β=1β=1 β=0β=0 β=1β=1 FedProx (Li et al., 2020) radius, Rμ≤R_μ≤ 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 α) −- −- −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 FedAvg (McMahan et al., 2017) reference 0.00.0 0.00.0 FedExP (Jhunjhunwala et al., 2023) adaptive extrapolation, ηg≥1 _g≥ 1 == ++ −0.4-0.4 +14.6+14.6 DiLoCo-style (Douillard et al., 2024) server step, outer momentum ++ ++ +2.0+2.0 +1.5+1.5 FedAdam (Reddi et al., 2021) server step αs/τs _s/ _s ++ ++ +17.2+17.2 +0.6+0.6 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. Δ is the mean percentage change across the 2424 configurations at each β. Results and Discussion. Across all 7272 configurations and all three random seeds, the Pearson correlation between the expected performance direction (encoded as −1-1, 00, and +1+1 for lower, similar, and higher perplexity, respectively) and the observed log-perplexity change relative to FedAvg is 0.710.71, increasing to 0.810.81 when restricted to β=1.0β=1.0. For the two heterogeneity regimes β∈0,1β∈\0,1\, the percentage changes in validation perplexity relative to FedAvg reported in Table 2, averaged over the 2424 configurations at each β and all three seeds, agree with 1313 of the 1616 expected lower/similar/higher outcomes. FedProx’s improvement over FedAvg is consistent with the potential benefit of tighter SSM envelopes identified in Theorem B.6 in Appendix B. SCAFFOLD’s improvement in 23/2423/24 configurations at β=1β=1, compared with deterioration in 22/2422/24 at β=0β=0, aligns with the disagreement-correction interpretation of the FedAvg bound developed in Appendix E. Together, these results support using the theorem’s displacement and disagreement terms to interpret the observed dependence of algorithm performance on client heterogeneity. 6 Conclusion In this work, we studied distributed learning with selective SSMs through architecture-aware theory, deriving bounds that connect federated optimization to recurrent stability, input-dependent discretization, projection norms, client drift, and heterogeneity. We then validated the derived bounds and provided experiments that demonstrated their relevance and practical applications. AI use statement In this work, we used generative AI tools for: assisting in the writing of proofs, formulating mathematical claims (specifically, refining author-derived mathematical statements in response to feedback on proof correctness), proposing or refining hypotheses, designing or providing feedback on research methodology or experiments, implementing methods, supporting qualitative and thematic data analysis and interpreting results. We have not used generative AI tools for: helping develop theoretical models or conceptual frameworks, generating synthetic data sets, providing critical ingredients for proving mathematical claims, assisting with translation, or cleaning and reformatting a dataset. Additionally, we used generative AI tools for: creating or modifying scientific figures or images (limited to suggestions for figure titles and labels), editing a research paper to improve readability, and suggesting a structure for a research paper. Our use of generative AI was three-fold. First, generative AI tools were used to provide feedback on theoretical arguments, hypotheses, methodological choices, and interpretations of experimental results. For all proofs, theorems, corollaries, and other mathematical statements in this work, the underlying claims and arguments were first derived and formalized. The resulting statements and proofs were subsequently provided to an AI agent for feedback on their correctness, logical consistency, and the validity of the stated claims. The agent’s feedback was independently evaluated, and any resulting revisions were made manually. Similarly, generative AI tools were used to check the consistency of our experimental analyses and interpretations with the reported results. All final analyses, interpretations, and claims were determined by the authors. Second, AI agents assisted with the implementation of the code used in our experiments. While some experiments use synthetic data generated by deterministic code, AI agents were only utilized to assist with writing the code, not generate the synthetic data. All AI-assisted code was manually inspected for correctness and subsequently evaluated using validation tests before being used in the reported experiments. Third, we used large language models to improve the presentation of the manuscript, including wording, readability, organization, section flow, and textual elements such as figure titles. The scientific figures and images themselves were not generated or modified using generative AI (though their titles and labels were adjusted with suggestions from AI models). All suggested edits were reviewed by the authors and incorporated manually where appropriate. We have reviewed all AI-assisted work and take responsibility for the final content of this work, including text, claims, code, and other artifacts produced with the aid of generative AI. Reproducibility Statement We provide complete proofs for the theoretical results presented in this work. For the empirical results, Appendix G provides detailed descriptions of the experimental setups, including the datasets and algorithms used, hyperparameter tuning procedures, and hardware used to conduct the experiments. We additionally provide our codebase containing implementations of the methods used to obtain the reported results, together with instructions for running the experiments as supplementary material. References Acar et al. (2021) D. A. E. Acar, Y. Zhao, R. Matas, M. Mattina, P. Whatmough, and V. Saligrama Federated learning based on dynamic regularization. In International Conference on Learning Representations, External Links: Link Cited by: Appendix E, 6th item, §2, Table 2. Agarwal et al. (2025) S. Agarwal, L. Ahmad, J. Ai, S. Altman, A. Applebaum, E. Arbus, R. K. Arora, Y. Bai, B. Baker, H. Bao, et al. Gpt-oss-120b & gpt-oss-20b model card. arXiv preprint arXiv:2508.10925. Cited by: §1. Blakeman et al. (2025) A. Blakeman, A. Basant, A. Khattar, A. Renduchintala, A. Bercovich, A. Ficek, A. Bjorlin, A. Taghibakhshi, A. S. Deshmukh, A. S. Mahabaleshwarkar, et al. Nemotron-h: a family of accurate and efficient hybrid mamba-transformer models. arXiv preprint arXiv:2504.03624. Cited by: §1, §2. Child et al. (2019) R. Child, S. Gray, A. Radford, and I. Sutskever Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. Cited by: §1. Choromanski et al. (2021) K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, D. B. Belanger, L. J. Colwell, and A. Weller Rethinking attention with performers. In International Conference on Learning Representations, External Links: Link Cited by: §1. Dao and Gu (2024) T. Dao and A. Gu Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning, External Links: Link Cited by: §2, §5. Douillard et al. (2024) A. Douillard, Q. Feng, A. A. Rusu, R. Chhaparia, Y. Donchev, A. Kuncoro, M. Ranzato, A. Szlam, and J. Shen DiLoCo: distributed low-communication training of language models. In 2nd Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ICML 2024), External Links: Link Cited by: Appendix E, 9th item, Table 2. Fang et al. (2026) W. Fang, D. Han, L. Yuan, S. Hosseinalipour, and C. G. Brinton Federated sketching lora: a flexible framework for heterogeneous collaborative fine-tuning of llms. International Conference on Machine Learning (ICML). Cited by: §1. Glorioso et al. (2024) P. Glorioso, Q. Anthony, Y. Tokpanov, J. Whittington, J. Pilault, A. Ibrahim, and B. Millidge Zamba: a compact 7b ssm hybrid model. arXiv preprint arXiv:2405.16712. Cited by: §2. Gu et al. (2020) A. Gu, T. Dao, S. Ermon, A. Rudra, and C. Ré HiPPO: recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (Eds.), Vol. 33, p. 1474–1487. External Links: Link Cited by: §2. Gu and Dao (2024) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, External Links: Link Cited by: §1, §1, §2. Gu et al. (2022a) A. Gu, K. Goel, A. Gupta, and C. Ré On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (Eds.), Vol. 35, p. 35971–35983. External Links: Link Cited by: §2. Gu et al. (2022b) A. Gu, K. Goel, and C. Re Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, External Links: Link Cited by: §1, §1, §2, §3. Hasani et al. (2023) R. Hasani, M. Lechner, T. Wang, M. Chahine, A. Amini, and D. Rus Liquid structural state-space models. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §2. He et al. (2025) W. He, K. Han, Y. Tang, C. Wang, Y. Yang, T. Guo, and Y. Wang DenseSSM: state space models with dense hidden connection for efficient large language models. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), L. Chiruzzo, A. Ritter, and L. Wang (Eds.), Albuquerque, New Mexico, p. 9243–9254. External Links: Link, Document, ISBN 979-8-89176-189-6 Cited by: §2. Hwang et al. (2024) S. Hwang, A. Lahoti, R. Puduppully, T. Dao, and A. Gu Hydra: bidirectional state space models through generalized matrix mixers. In Advances in Neural Information Processing Systems, A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (Eds.), Vol. 37, p. 110876–110908. External Links: Document, Link Cited by: §2. Ilharco et al. (2023) G. Ilharco, M. T. Ribeiro, M. Wortsman, L. Schmidt, H. Hajishirzi, and A. Farhadi Editing models with task arithmetic. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §F.1, §F.2. Jhunjhunwala et al. (2023) D. Jhunjhunwala, S. Wang, and G. Joshi FedExP: speeding up federated averaging via extrapolation. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: Appendix E, 8th item, Table 2. Kalman (1960) R. E. Kalman A new approach to linear filtering and prediction problems. Journal of Basic Engineering 82 (1), p. 35–45. External Links: ISSN 0021-9223 Cited by: §2. Karimireddy et al. (2020) S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh SCAFFOLD: stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning, H. D. I and A. Singh (Eds.), Proceedings of Machine Learning Research, Vol. 119, p. 5132–5143. External Links: Link Cited by: Appendix E, 3rd item, §2, Table 2. Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret Transformers are RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, H. D. I and A. Singh (Eds.), Proceedings of Machine Learning Research, Vol. 119, p. 5156–5165. External Links: Link Cited by: §1. Keçeci et al. (2025) E. Keçeci, M. Güzelkaya, and T. Kumbasar FedAlign: federated learning with state alignment for system identification. In 2025 International Conference on Control, Automation and Diagnosis (ICCAD), Vol. , p. 1–6. External Links: Document Cited by: Appendix E, 5th item, §1, §2, Table 2. Kim et al. (2026) H. Kim, B. Ko, M. Kang, M. Kim, D. Lee, J. Lee, S. Yoon, and D. Jung SSDi8: accurate and efficient 8-bit quantization for state space duality. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §2. Lahoti et al. (2026) A. Lahoti, K. Li, B. Chen, C. Wang, A. Bick, J. Z. Kolter, T. Dao, and A. Gu Mamba-3: improved sequence modeling using state space principles. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §1, §2, §3. Lehmann et al. (2015) J. Lehmann, R. Isele, M. Jakob, A. Jentzsch, D. Kontokostas, P. N. Mendes, S. Hellmann, M. Morsey, P. Van Kleef, S. Auer, et al. Dbpedia–a large-scale, multilingual knowledge base extracted from wikipedia. Semantic web 6 (2), p. 167–195. Cited by: §G.1. Li et al. (2021) Q. Li, B. He, and D. Song Model-contrastive federated learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 10713–10722. Cited by: §2. Li et al. (2024) S. Li, H. Singh, and A. Grover Mamba-nd: selective state space modeling for multi-dimensional data. In European Conference on Computer Vision, p. 75–92. Cited by: §2. Li et al. (2020) T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems, I. Dhillon, D. Papailiopoulos, and V. Sze (Eds.), Vol. 2, p. 429–450. External Links: Link Cited by: §D.1, 2nd item, §2, §3, Table 2. Li and Roth (2002) X. Li and D. Roth Learning question classifiers. In COLING 2002: The 19th International Conference on Computational Linguistics, Cited by: §G.1. Liu et al. (2025) T. Y. Liu, A. Achille, M. Trager, A. Golatkar, L. Zancato, and S. Soatto PICASO: permutation-invariant context composition with state space models. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §2. Ma et al. (2025) M. Ma, G. Liu, C. Cao, P. Deng, T. Dao, A. Gu, P. Jin, Z. Yang, Y. Xia, R. Luo, P. Hu, Z. Wang, Y. Chen, H. Liu, and T. Qin HybriDNA: a hybrid transformer-mamba2 long-range DNA language model. In ICLR 2025 Workshop on Machine Learning for Genomics Explorations, External Links: Link Cited by: §2. Maas et al. (2011) A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, p. 142–150. Cited by: §G.1. McMahan et al. (2017) B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, p. 1273–1282. Cited by: §D.1, 1st item, §1, §2, §3, Table 2. Peng et al. (2023) B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, X. Du, M. Grella, K. Gv, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. Wind, S. Woźniak, Z. Zhang, Q. Zhou, J. Zhu, and R. Zhu RWKV: reinventing RNNs for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, H. Bouamor, J. Pino, and K. Bali (Eds.), Singapore, p. 14048–14077. External Links: Link, Document Cited by: §1. Pióro et al. (2024) M. Pióro, M. Wołczyk, R. Pascanu, J. von Oswald, and J. Sacramento State soup: in-context skill learning, retrieval and mixing. arXiv preprint arXiv:2406.08423. Cited by: §2. Qiu et al. (2026) Z. Qiu, Z. Wang, X. Li, Y. Li, Y. Xu, Y. Wang, H. Zhang, R. Men, B. Mao, C. Zhang, et al. On the design of qwen3. 8-next architecture: evaluation, efficiency, and training stability. arXiv preprint arXiv:2608.30320. Cited by: §1. Qu et al. (2022) Z. Qu, X. Li, R. Duan, Y. Liu, B. Tang, and Z. Lu Generalized federated learning via sharpness aware minimization. In Proceedings of the 39th International Conference on Machine Learning, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato (Eds.), Proceedings of Machine Learning Research, Vol. 162, p. 18250–18280. External Links: Link Cited by: Appendix E, 7th item, Table 2. Reddi et al. (2021) S. J. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečný, S. Kumar, and H. B. McMahan Adaptive federated optimization. In International Conference on Learning Representations, External Links: Link Cited by: Appendix E, 4th item, §2, Table 2. Ren et al. (2025) L. Ren, Y. Liu, Y. Lu, yelong shen, C. Liang, and W. Chen Samba: simple hybrid state space models for efficient unlimited context language modeling. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §2. Soro et al. (2026) B. Soro, A. S. Zhang, B. Andreis, J. Jo, S. Chong, and S. J. Hwang LS-merge: merging language models in latent space. In The Fourteenth International Conference on Learning Representations, Cited by: §1. Tay et al. (2022) Y. Tay, M. Dehghani, D. Bahri, and D. Metzler Efficient transformers: a survey. ACM Computing Surveys 55 (6), p. 1–28. Cited by: §1. Team et al. (2026a) G. Team, S. E. Abd, V. Aggarwal, R. Algayres, A. Andreev, O. Bachem, I. Ballantyne, C. Brick, V. Cărbune, M. Casbon, et al. Gemma 4 technical report. arXiv preprint arXiv:2607.02770. Cited by: §1. Team et al. (2024) J. Team, B. Lenz, A. Arazi, A. Bergman, A. Manevich, B. Peleg, B. Aviram, C. Almagor, C. Fridman, D. Padnos, et al. Jamba-1.5: hybrid transformer-mamba models at scale. arXiv preprint arXiv:2408.12570. Cited by: §1. Team et al. (2026b) K. Team, T. Bai, Y. Bai, Y. Bao, J. Cai, X. Cai, P. Cao, Y. Cao, Z. Chai, Y. Charles, et al. Kimi k3: open frontier intelligence. arXiv preprint arXiv:2607.24653. Cited by: §1. Together Computer (2023) RedPajama: an open source recipe to reproduce llama training dataset External Links: Link Cited by: §G.1, §5. Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin Attention is all you need. In Advances in Neural Information Processing Systems, I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Eds.), Vol. 30, p. . External Links: Link Cited by: §1. Waleffe et al. (2024) R. Waleffe, W. Byeon, D. Riach, B. Norick, V. Korthikanti, T. Dao, A. Gu, A. Hatamizadeh, S. Singh, D. Narayanan, G. Kulshreshtha, V. Singh, J. Casper, J. Kautz, M. Shoeybi, and B. Catanzaro An empirical study of mamba-based language models. CoRR abs/2406.07887. External Links: Link Cited by: §2. Wang et al. (2020) J. Wang, Q. Liu, H. Liang, G. Joshi, and H. V. Poor Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems 33, p. 7611–7623. Cited by: §B.2, §B.2, §B.2, §2. Wortsman et al. (2022) M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, and L. Schmidt Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In Proceedings of the 39th International Conference on Machine Learning, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato (Eds.), Proceedings of Machine Learning Research, Vol. 162, p. 23965–23998. External Links: Link Cited by: §F.1. Xie et al. (2025) F. Xie, Z. Wang, W. Zhang, and C. Ma PVMamba: parallelizing vision mamba via dynamic state aggregation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 10218–10228. Cited by: §2. Yadav et al. (2023) P. Yadav, D. Tam, L. Choshen, C. A. Raffel, and M. Bansal TIES-merging: resolving interference when merging models. In Advances in Neural Information Processing Systems, A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (Eds.), Vol. 36, p. 7093–7115. External Links: Link Cited by: §F.2. Zaheer et al. (2020) M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. Big bird: transformers for longer sequences. Advances in neural information processing systems 33, p. 17283–17297. Cited by: §1. Zhang et al. (2026) W. Zhang, D. Deng, X. Wu, T. Zhang, X. Zheng, D. Niyato, and D. I. Kim Mitigating catastrophic forgetting in personalized federated learning for edge devices using state-space models. IEEE Transactions on Mobile Computing 25 (3), p. 3568–3582. External Links: Document Cited by: §1, §2. Zhang et al. (2015) X. Zhang, J. Zhao, and Y. LeCun Character-level convolutional networks for text classification. In Advances in Neural Information Processing Systems, C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (Eds.), Vol. 28, p. . External Links: Link Cited by: §G.1. Zhu et al. (2024) L. Zhu, B. Liao, Q. Zhang, X. Wang, W. Liu, and X. Wang Vision mamba: efficient visual representation learning with bidirectional state space model. In Forty-first International Conference on Machine Learning, External Links: Link Cited by: §2. Zuo et al. (2024) J. Zuo, M. Velikanov, D. E. Rhaiem, I. Chahed, Y. Belkada, G. Kunsch, and H. Hacid Falcon mamba: the first competitive attention-free 7b language model. arXiv preprint arXiv:2410.05355. Cited by: §1. Appendix Appendix A Selective SSM Gradient Bound Proof In this section, we derive an upper bound on the gradient norm of the loss function for some admissible input sequence processed by a single-layer selective state space model (SSM). Fix a parameter tuple W∈γ,δW _γ,δ satisfying the assumptions of Theorem 4.4. Let S≥1S≥ 1 be an arbitrary integer and ξ=(x,y∗)∈ξ=(x,y^*) an arbitrary admissible sequence pair of length S. We use the SSM dynamics from Section 3 and the notational conventions from section 4.1. Given these definitions, we examine the sensitivity of the loss with respect to the learnable model parameter tuple W:=(WΔ,WB,WC,A).W:=(W_ ,W_B,W_C,A). (12) Throughout the proof, gradients with respect to W are understood as block gradients, i.e., ‖∇Wf‖F:=(‖∇WΔf‖F2+‖∇WBf‖F2+‖∇WCf‖F2+‖∇Af‖F2)1/2.\| _Wf\|_F:= (\| _W_ f\|_F^2+\| _W_Bf\|_F^2+\| _W_Cf\|_F^2+\| _Af\|_F^2 )^1/2. (13) Next, let θ=vec⁡(W)θ=vec(W) concatenate the vectorized parameter blocks. Scalar gradients are column vectors after vectorization: ∇θf=vec⁡(∇Wf) _θf=vec ( _Wf ) For a vector- or matrix-valued function F, ∂F/∂W∂ F/∂ W denotes its flattened block Jacobian, with output entries indexing rows and scalar parameters indexing columns: ‖∂F∂W‖F2:=∑α‖∂F∂θα‖F2 \| ∂ F∂ W \|_F^2:= _α \| ∂ F∂ _α \|_F^2 In product-rule expressions, products and transposes involving matrix-valued derivatives are applied separately to each scalar-parameter slice before flattening. Since hth_t does not depend on WCW_C, its WCW_C-Jacobian block is zero. For clarity, in this section we omit the client index k and focus on bounding the gradient norm for a single client processing a sequence of length S. For the fixed sample ξ, using the simplified notation introduced in Section 3, the sequence loss is: ℒS(W)=1S∑t=1Sℓt(yt)L^S(W)= 1S _t=1^S _t(y_t) where ℓt(yt):=ℓt(yt(W,x),yt∗) _t(y_t):= _t(y_t(W;x),y^*_t) represents the instantaneous loss between the predicted and target tokens at the t-th place in the sequence S, and we omit the dependence on a fixed sample ξ. The goal of this derivation is to bound the norm of the gradient ∇WℒS(W) _WL^S(W). By the linearity of the gradient operator, it then follows that: ∇WℒS(W)=1S∑t=1S∇Wℓt(yt) _WL^S(W)= 1S _t=1^S _W _t(y_t) (14) We then apply the chain rule to the token loss ℓt(yt) _t(y_t), where ∇yℓt(yt) _y _t(y_t) is a column vector. Noting that yty_t depends on W both directly through the output projection vector CtC_t and indirectly through the hidden state hth_t, we obtain: ∇θℓt(yt) _θ _t(y_t) =(∂yt∂W)⊤∇yℓt(yt) = ( ∂ y_t∂ W ) _y _t(y_t) =(∂(ht⊤Ct)∂W)⊤∇yℓt(yt) = ( ∂(h_t C_t)∂ W ) _y _t(y_t) =(∂ht⊤∂WCt+ht⊤∂Ct∂W)⊤∇yℓt(yt). = ( ∂ h_t ∂ WC_t+h_t ∂ C_t∂ W ) _y _t(y_t). (15) Here, the product-rule terms inside the parentheses are formed separately for each scalar-parameter slice and constructed as the columns of the output Jacobian before taking the outer transpose. The first term contributes only to the WΔW_ , WBW_B, and A blocks, while the second term contributes only to the WCW_C block. Under the described block-gradient convention, the schematic chain-rule decomposition gives: ‖∇Wℓt(yt)‖F=‖vec⁡(∇Wℓt(yt))‖2=‖∇θℓt(yt)‖2=‖(∂ht⊤∂WCt+ht⊤∂Ct∂W)⊤∇yℓt(yt)‖2.\| _W _t(y_t)\|_F=\|vec( _W _t(y_t))\|_2=\| _θ _t(y_t)\|_2= \| ( ∂ h_t ∂ WC_t+h_t ∂ C_t∂ W ) _y _t(y_t) \|_2. (16) Next, applying the matrix-vector norm inequality to each scalar-parameter slice and summing the squared bounds gives: ‖∂ht⊤∂WCt‖F≤‖∂ht∂W‖F‖Ct‖2. \| ∂ h_t ∂ WC_t \|_F≤ \| ∂ h_t∂ W \|_F\|C_t\|_2. (17) Using this bound, the triangle inequality, and the fact that ‖M⊤v‖2≤‖M‖F‖v‖2\|M v\|_2≤\|M\|_F\|v\|_2, we get: ‖∇Wℓt(yt)‖F \| _W _t(y_t)\|_F ≤‖(∂ht⊤∂WCt)⊤∇yℓt(yt)‖2+‖(∂Ct∂W)⊤ht∇yℓt(yt)‖2 ≤ \| ( ∂ h_t ∂ WC_t ) _y _t(y_t) \|_2+ \| ( ∂ C_t∂ W ) h_t _y _t(y_t) \|_2 ≤‖∇yℓt(yt)‖2‖∂ht∂W‖F‖Ct‖2+‖(∂Ct∂W)⊤ht∇yℓt(yt)‖2. ≤\| _y _t(y_t)\|_2 \| ∂ h_t∂ W \|_F\|C_t\|_2+ \| ( ∂ C_t∂ W ) h_t _y _t(y_t) \|_2. (18) Next, utilizing Assumptions 4.2 and 4.3, we bound the two terms of this expression separately. First, since Ct=WCxtC_t=W_Cx_t, we have: ‖Ct‖2≤‖WC‖2‖xt‖2≤‖WC‖2Lx\|C_t\|_2≤\|W_C\|_2\|x_t\|_2≤\|W_C\|_2L_x (19) Therefore, ‖∇yℓt(yt)‖2‖∂ht∂W‖F‖Ct‖2≤LℓLx‖WC‖2‖∂ht∂W‖F \| _y _t(y_t) \|_2 \| ∂ h_t∂ W \|_F\|C_t\|_2≤ L_ L_x\|W_C\|_2 \| ∂ h_t∂ W \|_F (20) For the second term, since hth_t does not depend on WCW_C and xtx_t is fixed, equation (2b) gives: ∇WCℓt(yt)=ht∇yℓt(yt)xt⊤. _W_C _t(y_t)=h_t _y _t(y_t)x_t . (21) Furthermore, since Ct=WCxtC_t=W_Cx_t, the Jacobian ∂Ct/∂W∂ C_t/∂ W is zero in the WΔW_ , WBW_B, and A parameter blocks, and nonzero only for the WCW_C block. As a result, (∂Ct/∂W)⊤ht∇yℓt(yt)(∂ C_t/∂ W) h_t _y _t(y_t) contains the vectorized WCW_C-block gradient and zeros in all other parameter blocks, yielding: ‖(∂Ct∂W)⊤ht∇yℓt(yt)‖2 \| ( ∂ C_t∂ W ) h_t _y _t(y_t) \|_2 =‖∇WCℓt(yt)‖F=‖ht∇yℓt(yt)xt⊤‖F. =\| _W_C _t(y_t)\|_F=\|h_t _y _t(y_t)x_t \|_F. (22) This can be bounded as: ‖ht∇yℓt(yt)xt⊤‖F \|h_t _y _t(y_t)x_t \|_F =‖ht∇yℓt(yt)‖2‖xt‖2≤‖ht∇yℓt(yt)‖2Lx = \|h_t _y _t(y_t) \|_2\|x_t\|_2≤ \|h_t _y _t(y_t) \|_2L_x (23) ≤‖ht‖F‖∇yℓt(yt)‖2Lx≤LxLℓ‖ht‖F ≤ \|h_t \|_F \| _y _t(y_t) \|_2L_x≤ L_xL_ \|h_t \|_F (24) By combining this result with (20), we can rewrite inequality (18) as: ‖∇Wℓt(yt)‖F≤LℓLx(‖ht‖F+‖WC‖2‖∂ht∂W‖F) \| _W _t(y_t) \|_F≤ L_ L_x ( \|h_t \|_F+ \|W_C \|_2 \| ∂ h_t∂ W \|_F ) (25) Next, for the remaining two terms, we first evaluate the bound on the hidden state matrix ht∈ℝN×Dh_t ^N× D. Given an initialization of h0=h_0=0 (which is standard when processing independent training sequences), unrolling the recurrence from equation (2a) yields: ht=∑j=1t(⨀k=j+1tA¯k)⊙(B¯jdiag⁡(xj))h_t= _j=1^t ( _k=j+1^t A_k ) ( B_jdiag(x_j) ) (26) where we define the empty cumulative Hadamard product as the all-ones matrix, ⨀k=t+1tA¯k=N×D _k=t+1^t A_k=1_N× D. Then, taking the Frobenius norm of both sides and applying the triangle inequality, we get: ‖ht‖F≤∑j=1t‖(⨀k=j+1tA¯k)⊙(B¯jdiag⁡(xj))‖F=∑j=1t‖exp⁡(A∑k=j+1tΔk⊤)⊙(B¯jdiag⁡(xj))‖F\|h_t\|_F≤ _j=1^t \| ( _k=j+1^t A_k ) ( B_jdiag(x_j) ) \|_F= _j=1^t \| (A _k=j+1^t _k ) ( B_jdiag(x_j) ) \|_F (27) Utilizing the property of Hadamard product with Frobenius norms, the fact that ‖B¯jdiag⁡(xj)‖F≤‖B¯j‖2‖diag⁡(xj)‖F≤‖B¯j‖2Lx \| B_jdiag(x_j) \|_F≤ \| B_j \|_2 \|diag(x_j) \|_F≤ \| B_j \|_2L_x, and the definition of B¯t B_t, it follows that: ‖exp⁡(A∑k=j+1tΔk⊤)⊙(B¯jdiag⁡(xj))‖F≤‖exp⁡(A∑k=j+1tΔk⊤)‖F‖WBxjΔj⊤‖2Lx \| (A _k=j+1^t _k ) ( B_jdiag(x_j) ) \|_F≤ \| (A _k=j+1^t _k ) \|_F \|W_Bx_j _j \|_2L_x (28) Then: ‖WBxjΔj⊤‖2Lx≤‖WB‖2‖xj‖2‖Δj⊤‖2Lx≤‖WB‖2‖Softplus(WΔxj)‖2Lx2 \|W_Bx_j _j \|_2L_x≤ \|W_B \|_2 \|x_j \|_2 \| _j \|_2L_x≤ \|W_B \|_2 \|Softplus (W_ x_j ) \|_2L_x^2 (29) And since Softplus(a)≤|a|+log⁡(2)Softplus(a)≤|a|+ (2) so ‖Softplus()‖2≤‖||+log⁡(2)‖2≤‖2+Dlog⁡(2) \|Softplus(x) \|_2≤ \||x|+ (2)1 \|_2≤\|x\|_2+ D (2), it follows that ‖Softplus(WΔxj)‖2≤‖|WΔxj|+log⁡(2)1‖2 \|Softplus (W_ x_j ) \|_2≤ \| |W_ x_j |+ (2)1 \|_2, and from the triangle inequality, this results in ‖|WΔxj|+log⁡(2)1‖2≤‖WΔ‖2Lx+Dlog⁡(2) \| |W_ x_j |+ (2)1 \|_2≤ \|W_ \|_2L_x+ D (2) So the bound becomes: ‖WBxjΔj⊤‖2Lx≤‖WB‖2Lx2(‖WΔ‖2Lx+Dlog⁡(2)) \|W_Bx_j _j \|_2L_x≤ \|W_B \|_2L_x^2 ( \|W_ \|_2L_x+ D (2) ) (30) Since these terms are independent of j, it follows that: ‖ht‖F≤‖WB‖2Lx2(‖WΔ‖2Lx+Dlog⁡(2))∑j=1t‖exp⁡(A∑k=j+1tSoftplus(WΔxk)⊤)‖F\|h_t\|_F≤ \|W_B \|_2L_x^2 ( \|W_ \|_2L_x+ D (2) ) _j=1^t \| (A _k=j+1^tSoftplus (W_ x_k ) ) \|_F (31) Next, the term ∂ht∂W ∂ h_t∂ W can be expanded from equation (2a) as: ∂ht∂W=∂A¯t∂W⊙ht−1+A¯t⊙∂ht−1∂W+∂B¯t∂Wdiag⁡(xt) ∂ h_t∂ W= ∂ A_t∂ W h_t-1+ A_t ∂ h_t-1∂ W+ ∂ B_t∂ Wdiag(x_t) (32) We bound each term individually, using the fact that: ‖∂ht∂W‖F≤‖∂ht∂WΔ‖F+‖∂ht∂WB‖F+‖∂ht∂WC‖F+‖∂ht∂A‖F \| ∂ h_t∂ W \|_F≤ \| ∂ h_t∂ W_ \|_F+ \| ∂ h_t∂ W_B \|_F+ \| ∂ h_t∂ W_C \|_F+ \| ∂ h_t∂ A \|_F (33) where the WCW_C-dependent term is zero. By the triangle inequality: ‖∂ht∂W‖F≤‖∂A¯t∂W⊙ht−1‖F+‖A¯t⊙∂ht−1∂W‖F+‖∂B¯t∂Wdiag⁡(xt)‖F \| ∂ h_t∂ W \|_F≤ \| ∂ A_t∂ W h_t-1 \|_F+ \| A_t ∂ h_t-1∂ W \|_F+ \| ∂ B_t∂ Wdiag(x_t) \|_F (34) Then, we bound each term. First: ‖∂B¯t∂Wdiag⁡(xt)‖F \| ∂ B_t∂ Wdiag(x_t) \|_F ≤‖∂B¯t∂W‖F‖diag⁡(xt)‖F≤Lx(‖∂B¯t∂WB‖F+‖∂B¯t∂WΔ‖F) ≤ \| ∂ B_t∂ W \|_F \|diag(x_t) \|_F≤ L_x ( \| ∂ B_t∂ W_B \|_F+ \| ∂ B_t∂ W_ \|_F ) ≤Lx(N‖xt‖2‖Δt‖2+‖WBxt‖2‖∂Δt∂WΔ‖F) ≤ L_x ( N \|x_t \|_2 \| _t \|_2+ \|W_Bx_t \|_2 \| ∂ _t∂ W_ \|_F ) ≤Lx(NLx(‖WΔ‖2Lx+Dlog⁡(2))+D‖WB‖2Lx2), ≤ L_x ( NL_x ( \|W_ \|_2L_x+ D (2) )+ D \|W_B \|_2L_x^2 ), (35) where we used the fact that Softplus′⁡(z)=σ⁡(z)≤1Softplus (z)=σ(z)≤ 1, so since Δt=Softplus⁡(WΔxt) _t=Softplus(W_ x_t), it is true that ‖∂Δt∂WΔ‖F≤D‖xt‖2≤DLx. \| ∂ _t∂ W_ \|_F≤ D\|x_t\|_2≤ DL_x. Next: ‖∂A¯t∂W⊙ht−1‖F \| ∂ A_t∂ W h_t-1 \|_F ≤‖∂A¯t∂W‖F‖ht−1‖F≤(‖∂A¯t∂A‖F+‖∂A¯t∂WΔ‖F)‖ht−1‖F ≤ \| ∂ A_t∂ W \|_F \|h_t-1 \|_F≤ ( \| ∂ A_t∂ A \|_F+ \| ∂ A_t∂ W_ \|_F ) \|h_t-1 \|_F ≤(‖A¯t‖F‖Δt‖2+‖∂A¯t∂Δt‖F‖∂Δt∂WΔ‖F)‖ht−1‖F ≤ ( \| A_t \|_F \| _t \|_2+ \| ∂ A_t∂ _t \|_F \| ∂ _t∂ W_ \|_F ) \|h_t-1 \|_F ≤(‖A¯t‖F‖Δt‖2+D‖A‖2‖A¯t‖FLx)‖ht−1‖F ≤ ( \| A_t \|_F \| _t \|_2+ D \|A \|_2 \| A_t \|_FL_x ) \|h_t-1 \|_F ≤‖ht−1‖F‖A¯t‖F((‖WΔ‖2Lx+Dlog⁡(2))+D‖A‖2Lx) ≤ \|h_t-1 \|_F \| A_t \|_F ( ( \|W_ \|_2L_x+ D (2) )+ D \|A \|_2L_x ) (36) Finally, by definition of ‖M‖max:=maxi,j⁡|Mij|\|M\|_ := _i,j|M_ij|, it follows that: ‖A¯t⊙∂ht−1∂W‖F≤‖A¯t‖max‖∂ht−1∂W‖F \| A_t ∂ h_t-1∂ W \|_F≤ \| A_t \|_ \| ∂ h_t-1∂ W \|_F (37) Combining all the terms together, we get the following recurrence: ‖∂ht∂W‖F \| ∂ h_t∂ W \|_F ≤‖WB‖2Lx2(‖WΔ‖2Lx+Dlog⁡(2))∑j=1t−1‖exp⁡(A∑k=j+1t−1Softplus(WΔxk)⊤)‖F ≤ \|W_B \|_2L_x^2 ( \|W_ \|_2L_x+ D (2) ) _j=1^t-1 \| (A _k=j+1^t-1Softplus (W_ x_k ) ) \|_F ×‖A¯t‖F((‖WΔ‖2Lx+Dlog⁡(2))+D‖A‖2Lx) × \| A_t \|_F ( ( \|W_ \|_2L_x+ D (2) )+ D \|A \|_2L_x ) +Lx(NLx(‖WΔ‖2Lx+Dlog⁡(2))+D‖WB‖2Lx2)+‖A¯t‖max‖∂ht−1∂W‖F +L_x ( NL_x ( \|W_ \|_2L_x+ D (2) )+ D \|W_B \|_2L_x^2 )+ \| A_t \|_ \| ∂ h_t-1∂ W \|_F (38) Next, define: KΔ K_ :=‖WΔ‖2Lx+Dlog⁡(2) := \|W_ \|_2L_x+ D (2) (39) KB K_B :=‖WB‖2Lx2 := \|W_B \|_2L_x^2 (40) Φj,t _j,t :=‖exp⁡(A∑k=j+1tSoftplus(WΔxk)⊤)‖F := \| (A _k=j+1^tSoftplus (W_ x_k ) ) \|_F (41) So the recurrence becomes: ‖∂ht∂W‖F≤ \| ∂ h_t∂ W \|_F≤ ‖A¯t‖FKBKΔ(KΔ+D‖A‖2Lx)∑j=1t−1Φj,t−1+Lx2NKΔ+LxDKB \| A_t \|_FK_BK_ (K_ + D \|A \|_2L_x ) _j=1^t-1 _j,t-1+L_x^2 NK_ +L_x DK_B +‖A¯t‖max‖∂ht−1∂W‖F + \| A_t \|_ \| ∂ h_t-1∂ W \|_F (42) Since h0=h_0=0 is parameter-independent, ∂h0/∂W=0∂ h_0/∂ W=0, iterating the inequality gives: ‖∂ht∂W‖F≤∑m=1t[ \| ∂ h_t∂ W \|_F≤ _m=1^t [ ‖A¯m‖FKBKΔ(KΔ+D‖A‖2Lx)(∑j=1m−1Φj,m−1) \| A_m \|_FK_BK_ (K_ + D \|A \|_2L_x ) ( _j=1^m-1 _j,m-1 ) +Lx2NKΔ+LxDKB]∏n=m+1t‖A¯n‖max, +L_x^2 NK_ +L_x DK_B ] _n=m+1^t \| A_n \|_ , (43) where empty sums are calculated as zero and empty scalar products are calculated as one. Next, by Definition 4.1, each entry of AΔt⊤A _t is at most −γδ-γδ. Since A¯t=exp⁡(AΔt⊤) A_t= (A _t ) element-wise, it follows that: ‖A¯t‖max≤exp⁡(−γδ):=ρ<1 \| A_t \|_ ≤ (-γδ ):=ρ<1 (44) as well as ‖A¯t‖F≤NDρ \| A_t \|_F≤ NDρ (45) and in addition, by the definition of Φj,t _j,t, the following inequality holds: Φj,m−1≤NDρm−1−j _j,m-1≤ NDρ^m-1-j (46) As a result, the inner summation becomes a simple geometric series: ∑j=1m−1Φj,m−1≤ND∑j=1m−1ρm−1−j=ND∑r=0m−2ρr≤ND1−ρ _j=1^m-1 _j,m-1≤ ND _j=1^m-1ρ^m-1-j= ND _r=0^m-2ρ^r≤ ND1-ρ (47) As such, we can group all time-independent terms under a new constant: KW:=NDKBKΔ(KΔ+D‖A‖2Lx)(ρ1−ρ)+Lx2NKΔ+LxDKBK_W:=NDK_BK_ (K_ + D \|A \|_2L_x ) ( ρ1-ρ )+L_x^2 NK_ +L_x DK_B (48) reducing the recurrence to: ‖∂ht∂W‖F≤KW∑m=1tρt−m≤KW1−ρ \| ∂ h_t∂ W \|_F≤ K_W _m=1^tρ^t-m≤ K_W1-ρ (49) Finally, combining everything, we get the final bound on the gradient of loss at time t: ‖∇Wℓt(yt)‖F≤LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ) \| _W _t(y_t) \|_F≤ L_ L_x (K_BK_ ND1-ρ+ \|W_C \|_2 K_W1-ρ ) (50) Since ∇WℒS(W)=1S∑t=1S∇Wℓt(yt) _WL^S(W)= 1S _t=1^S _W _t(y_t), by the triangle inequality the same bound holds for the averaged empirical loss: ‖∇WℒS(W)‖F≤LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ)=SSM(W) \| _WL^S(W) \|_F≤ L_ L_x (K_BK_ ND1-ρ+ \|W_C \|_2 K_W1-ρ )=G_SSM(W) (51) Since S≥1S≥ 1 and the admissible length-S sequence pairs ξ∈ξ were arbitrary, and the right-hand side of the inequality is independent of S and ξ, the bound holds uniformly over all such lengths and admissible sequence pairs for the fixed parameter tuple W, proving Theorem 4.4. ∎ Appendix B Convergence of Federated SSM Baselines In this section, we first define and prove a set of lemmas required to instantiate standard federated optimization convergence results in terms of selective SSM parameters. We then provide the proof of Theorem B.5, and state the corresponding FedProx theorem. B.1 Auxiliary Lemmas for SSM Federated Convergence Rather than proving the FedAvg or FedProx convergence bounds directly, we first show how the generic smoothness, gradient second-moment, stochastic variance, and client heterogeneity quantities appearing in those bounds can be replaced by SSM-specific constants. In particular, we establish local control of the SSM gradient bound in Lemma B.1, prove the local smoothness constant LSSM,RL_SSM,R in Lemma 4.7, and derive the architecture-dependent quantities GSSM,R2G_SSM,R^2, σSSM,R,Ξ2 _SSM,R, ^2, and ΓSSM,R2 _SSM,R^2 in Lemmas B.2, B.3, and B.4, respectively. These lemmas are then used in the remainder of this section to state the final convergence bounds for FedAvg and FedProx distributed learning baselines. B.1.1 Local Lipschitz Continuity of the SSM Gradient Bound Lemma and Proof Lemma B.1 (Local Lipschitz Continuity of the SSM Gradient Bound). Let WrW^r be some fixed reference parameter tuple and let 0<R<∞0<R<∞ be a finite radius such that R(Wr)⊆γ,δN_R(W^r) _γ,δ. Suppose Assumptions 4.2 and 4.3 hold with the neighborhood-uniform constants specified in Section 4.2. Let SSM(W)G_SSM(W) be the bounding function defined in Theorem 4.4. Then SSMG_SSM is Lipschitz continuous on R(Wr)N_R(W^r). Specifically, with N,D,Lℓ,Lx,γ,δN,D,L_ ,L_x,γ,δ fixed, there exists a finite constant LSSM,R>0L_G_SSM,R>0, depending only on R and the block norms ‖WΔr‖2\|W_ ^r\|_2, ‖WBr‖2\|W_B^r\|_2, ‖WCr‖2\|W_C^r\|_2, and ‖Ar‖2\|A^r\|_2, such that for all W,V∈R(Wr)W,V _R(W^r), |SSM(W)−SSM(V)|≤LSSM,R‖W−V‖F.|G_SSM(W)-G_SSM(V)|≤ L_G_SSM,R\|W-V\|_F. (52) Proof of Lemma B.1: Fix the reference tuple WrW^r and the radius R. Throughout this proof, N,D,Lℓ,Lx,γ,δN,D,L_ ,L_x,γ,δ are constant as W and V vary over the neighborhood R(Wr)N_R(W^r). In particular, ρ=exp⁡(−γδ)ρ= (-γδ) is the same for both parameter tuples. First, by Theorem 4.4, SSM(W)G_SSM(W) has an explicit definition: SSM(W)=LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ),G_SSM(W)=L_ L_x (K_BK_ ND1-ρ+ \|W_C \|_2 K_W1-ρ ), (53) where W is the tuple of selective SSM weights (WΔ,WB,WC,A)(W_ ,W_B,W_C,A). Since W=(WΔ,WB,WC,A)W=(W_ ,W_B,W_C,A) is a parameter tuple, given another parameter tuple V=(VΔ,VB,VC,VA)V=(V_ ,V_B,V_C,V_A) we use the block Frobenius norm: ‖W−V‖F2:=‖WΔ−VΔ‖F2+‖WB−VB‖F2+‖WC−VC‖F2+‖A−VA‖F2.\|W-V\|_F^2:=\|W_ -V_ \|_F^2+\|W_B-V_B\|_F^2+\|W_C-V_C\|_F^2+\|A-V_A\|_F^2. (54) For each parameter block X∈Δ,B,C,AX∈\ ,B,C,A\, let WXW_X denote the corresponding block of W, and let WXrW_X^r denote the corresponding block of the central model WrW^r (for notational uniformity, we write WA:=AW_A:=A and WAr:=ArW_A^r:=A^r). Given the central model WrW^r, we evaluate the maximum possible value of a norm of each parameter matrix within the neighborhood R(Wr)N_R(W^r). Since W∈R(Wr)W _R(W^r), each block satisfies ‖WX−WXr‖F≤‖W−Wr‖F≤R.\|W_X-W_X^r\|_F≤\|W-W^r\|_F≤ R. (55) Therefore, by the triangle inequality and the fact that ∥⋅∥2≤∥⋅∥F\|·\|_2≤\|·\|_F, ‖WX‖2≤‖WXr‖2+‖WX−WXr‖2≤‖WXr‖2+‖WX−WXr‖F≤‖WXr‖2+R.\|W_X\|_2≤\|W_X^r\|_2+\|W_X-W_X^r\|_2≤\|W_X^r\|_2+\|W_X-W_X^r\|_F≤\|W_X^r\|_2+R. (56) Next, define: MΔ M_ :=‖WΔr‖2+R := \|W_ ^r \|_2+R (57) MB M_B :=‖WBr‖2+R := \|W_B^r \|_2+R (58) MC M_C :=‖WCr‖2+R := \|W_C^r \|_2+R (59) MA M_A :=‖Ar‖2+R := \|A^r \|_2+R (60) Then, by (56), it is true that for all individual matrices WB,WC,WΔ,AW_B,W_C,W_ ,A, the following holds: ‖WΔ‖2 \|W_ \|_2 ≤MΔ ≤ M_ (61) ‖WB‖2 \|W_B \|_2 ≤MB ≤ M_B (62) ‖WC‖2 \|W_C \|_2 ≤MC ≤ M_C (63) ‖A‖2 \|A \|_2 ≤MA ≤ M_A (64) The same bounds also hold for the corresponding blocks of any V∈R(Wr)V _R(W^r). We further define: MKΔ M_K_ :=MΔLx+Dlog⁡(2) :=M_ L_x+ D (2) (65) MKB M_K_B :=MBLx2 :=M_BL_x^2 (66) Next, we evaluate the Lipschitz constants for individual components of the expression for SSM(W)G_SSM(W). We’l utilize the property that: |‖X‖2−‖Y‖2|≤‖X−Y‖2≤‖X−Y‖F | \|X \|_2- \|Y \|_2 |≤ \|X-Y \|_2≤ \|X-Y \|_F. First: |KΔ(W)−KΔ(V)| |K_ (W)-K_ (V) | ≤Lx|‖WΔ‖2−‖VΔ‖2|≤Lx‖WΔ−VΔ‖F≤Lx‖W−V‖F ≤ L_x | \|W_ \|_2- \|V_ \|_2 |≤ L_x \|W_ -V_ \|_F≤ L_x \|W-V \|_F (67) |KB(W)−KB(V)| |K_B(W)-K_B(V) | ≤Lx2|‖WB‖2−‖VB‖2|≤Lx2‖WB−VB‖F≤Lx2‖W−V‖F ≤ L_x^2 | \|W_B \|_2- \|V_B \|_2 |≤ L_x^2 \|W_B-V_B \|_F≤ L_x^2 \|W-V \|_F (68) So the Lipschitz constants for the base functions are: LKΔ L_K_ =Lx =L_x (69) LKB L_K_B =Lx2 =L_x^2 (70) LC L_C =1 (for the term ‖WC‖2) =1 (for the term \|W_C \|_2) (71) LA L_A =1 (for the term ‖A‖2) =1 (for the term \|A \|_2) (72) Next, by the product rule for Lipschitz functions, if |f|≤Mf|f|≤ M_f and |g|≤Mg|g|≤ M_g and are Lipschitz with constants LfL_f and LgL_g on R(Wr)N_R(W^r), then fgfg is Lipschitz with constant Lfg=MfLg+MgLf.L_fg=M_fL_g+M_gL_f. (73) Given that, let P1(W):=KB(W)KΔ(W).P_1(W):=K_B(W)K_ (W). (74) Then LP1=MKBLKΔ+MKΔLKB=MKBLx+MKΔLx2,L_P_1=M_K_BL_K_ +M_K_ L_K_B=M_K_BL_x+M_K_ L_x^2, (75) and MP1=MKBMKΔ.M_P_1=M_K_BM_K_ . (76) Next, define P2(W):=KΔ(W)+D‖A‖2Lx.P_2(W):=K_ (W)+ D\|A\|_2L_x. (77) Then, by the sum rule of Lipschitz functions: LP2=LKΔ+DLALx=Lx+DLx,L_P_2=L_K_ + DL_AL_x=L_x+ DL_x, (78) and MP2=MKΔ+DMALx.M_P_2=M_K_ + DM_AL_x. (79) Continuing, define P3(W):=P1(W)P2(W).P_3(W):=P_1(W)P_2(W). (80) Then LP3=MP1LP2+MP2LP1,L_P_3=M_P_1L_P_2+M_P_2L_P_1, (81) and MP3=MP1MP2.M_P_3=M_P_1M_P_2. (82) Given these definitions, and using the definition of KWK_W from Theorem 4.4, a valid Lipschitz constant for KWK_W is: LKW=NDρ1−ρLP3+Lx3N+Lx3D,L_K_W=ND ρ1-ρL_P_3+L_x^3 N+L_x^3 D, (83) and the upper bound for KW(W)K_W(W) is: MKW=NDρ1−ρMP3+Lx2NMKΔ+LxDMKB.M_K_W=ND ρ1-ρM_P_3+L_x^2 NM_K_ +L_x DM_K_B. (84) Finally, define P4(W):=‖WC‖2KW(W).P_4(W):=\|W_C\|_2K_W(W). (85) Since ‖WC‖2\|W_C\|_2 is bounded by MCM_C and has Lipschitz constant LC=1L_C=1, the product rule gives LP4=MCLKW+MKW.L_P_4=M_CL_K_W+M_K_W. (86) Combining everything together and using linearity, we get the final Lipschitz constant for SSM(W)G_SSM(W): LSSM,R=LℓLx(ND1−ρLP1+11−ρLP4).L_G_SSM,R=L_ L_x ( ND1-ρL_P_1+ 11-ρL_P_4 ). (87) Because R<∞R<∞, the block norms of WrW^r are finite, and ρ<1ρ<1, all constants MKΔM_K_ , MKBM_K_B, MP1M_P_1, MP2M_P_2, MP3M_P_3, MKWM_K_W, LP1L_P_1, LP2L_P_2, LP3L_P_3, LKWL_K_W, and LP4L_P_4 are finite. As such, LSSM,R<∞L_G_SSM,R<∞, and ∀W,V∈R(Wr)∀ W,V _R(W^r), |SSM(W)−SSM(V)|≤LSSM,R‖W−V‖F.|G_SSM(W)-G_SSM(V)|≤ L_G_SSM,R\|W-V\|_F. (88) This proves that SSM(W)G_SSM(W) is locally Lipschitz continuous on R(Wr)N_R(W^r). ∎ B.1.2 Local SSM smoothness Lemma Proof Proof of Lemma 4.7: Fix the reference tuple WrW^r and the radius R from the lemma. Let ξ=(x,y∗)∈ξ=(x,y^*) be an arbitrary admissible sequence pair of finite length S≥1S≥ 1. The inputs and targets are held fixed when differentiating with respect to the model parameters. For this fixed sequence pair, we use a simplified notation for ht(W,x)h_t(W;x), yt(W,x)y_t(W;x), and ℓt(yt(W,x),yt∗) _t(y_t(W;x),y_t^*) as ht(W)h_t(W), yt(W)y_t(W), and ℓt(yt(W)) _t(y_t(W)), respectively. Throughout the proof, Lℓ,Hℓ,Lx,γ,δL_ ,H_ ,L_x,γ,δ are fixed uniformly over R(Wr)N_R(W^r), with ρ=exp⁡(−γδ)<1ρ= (-γδ)<1. The estimates below hold for every W∈R(Wr)W _R(W^r), every finite sequence length S≥1S≥ 1, every admissible length-S pair ξ∈ξ , and every token position t≤St≤ S, with constants independent of S, ξ, and t. We use the parameter envelopes from (57)-(60) and the induced bounds from (66), (65), and (84). For convenience, we define: HR H_R :=MKBMKΔND1−ρ :=M_K_BM_K_ ND1-ρ (89) JR J_R :=MKW1−ρ := M_K_W1-ρ (90) By the definition of MXM_X as an upper-bound and from the Proof of Theorem 4.4, we know the following hold: ‖ht(W)‖F \|h_t(W) \|_F ≤KB(W)KΔ(W)ND1−ρ≤HR ≤ K_B(W)K_ (W) ND1-ρ≤ H_R (91) ‖∂ht(W)∂W‖F \| ∂ h_t(W)∂ W \|_F ≤KW(W)1−ρ≤JR ≤ K_W(W)1-ρ≤ J_R (92) Given this setup, we begin by bounding D2htD^2h_t where D is the Fréchet derivative defined: DF(W)[U] DF(W)[U] :=dϵF(W+ϵU)|ϵ=0 := ddεF(W+ε U) |_ε=0 D2F(W)[U,V] D^2F(W)[U,V] :=∂2∂ϵ∂ζF(W+ϵU+ζV)|ϵ=ζ=0 := ∂^2∂ε∂ζF(W+ε U+ζ V) |_ε=ζ=0 with U,VU,V being unit perturbations. We also utilize the operator norm defined ‖DΨ(W)‖op:=sup‖U‖F=1‖DΨ(W)[U]‖F\|D (W)\|_op:= _\|U\|_F=1\|D (W)[U]\|_F. Note that after flattening the input and output tensors, this corresponds to the usual spectral norm of the corresponding Jacobian matrix. Parameter Hessians are represented as matrices using the vectorization of parameter blocks introduced in the single-layer proof. Their operator norms are taken with respect to the corresponding Euclidean norm, which equals the block Frobenius norm before vectorization. Let U=(UΔ,UB,UC,UA)U=(U_ ,U_B,U_C,U_A) and V=(VΔ,VB,VC,VA)V=(V_ ,V_B,V_C,V_A) be arbitrary block perturbations satisfying ‖U‖F=‖V‖F=1.\|U\|_F=\|V\|_F=1. (93) Each perturbation block therefore has Frobenius norm at most one. Since Δt _t depends only on WΔW_ , we write DΔt(W)[UΔ]D _t(W)[U_ ] as shorthand for DΔt(W)[(UΔ,0,0,0)]D _t(W)[(U_ ,0,0,0)], and use the analogous convention for its second derivative. By the property of the operator norm (with Ψ(W)=ht(W) (W)=h_t(W)), ‖Dht(W)[U]‖F≤‖Dht(W)‖op‖U‖F=‖Dht(W)‖op≤‖∂ht(W)∂W‖F≤JR. \|Dh_t(W)[U] \|_F≤\|Dh_t(W)\|_op\|U\|_F=\|Dh_t(W)\|_op≤ \| ∂ h_t(W)∂ W \|_F≤ J_R. (94) The same argument gives ‖Dht(W)[V]‖F≤JR.\|Dh_t(W)[V]\|_F≤ J_R. (95) and therefore it follows that ‖Dht(W)[U]‖F≤JR,‖Dht(W)[V]‖F≤JR.\|Dh_t(W)[U]\|_F≤ J_R, \|Dh_t(W)[V]\|_F≤ J_R. (96) Next, let bt(W):=B¯t(W)diag⁡(xt)b_t(W):= B_t(W)diag(x_t), then the hidden state recurrence equation becomes: ht(W)=A¯t(W)⊙ht−1(W)+bt(W),h_t(W)= A_t(W) h_t-1(W)+b_t(W), (97) and taking the first Fréchet derivative in direction U gives: Dht(W)[U]=DA¯t(W)[U]⊙ht−1(W)+A¯t(W)⊙Dht−1(W)[U]+Dbt(W)[U].Dh_t(W)[U]=D A_t(W)[U] h_t-1(W)+ A_t(W) Dh_t-1(W)[U]+Db_t(W)[U]. (98) Taking a second Fréchet derivative in direction V gives D2ht(W)[U,V]= D^2h_t(W)[U,V]= D2A¯t(W)[U,V]⊙ht−1(W) \ D^2 A_t(W)[U,V] h_t-1(W) +DA¯t(W)[U]⊙Dht−1(W)[V] +D A_t(W)[U] Dh_t-1(W)[V] +DA¯t(W)[V]⊙Dht−1(W)[U] +D A_t(W)[V] Dh_t-1(W)[U] +A¯t(W)⊙D2ht−1(W)[U,V] + A_t(W) D^2h_t-1(W)[U,V] +D2bt(W)[U,V]. +D^2b_t(W)[U,V]. (99) Next, taking the Frobenius norms, and by the triangle inequality, it follows that ‖D2ht(W)[U,V]‖F≤ \|D^2h_t(W)[U,V]\|_F≤ ‖D2A¯t(W)[U,V]⊙ht−1(W)‖F \|D^2 A_t(W)[U,V] h_t-1(W)\|_F +‖DA¯t(W)[U]⊙Dht−1(W)[V]‖F +\|D A_t(W)[U] Dh_t-1(W)[V]\|_F +‖DA¯t(W)[V]⊙Dht−1(W)[U]‖F +\|D A_t(W)[V] Dh_t-1(W)[U]\|_F +‖A¯t(W)⊙D2ht−1(W)[U,V]‖F +\| A_t(W) D^2h_t-1(W)[U,V]\|_F +‖D2bt(W)[U,V]‖F. +\|D^2b_t(W)[U,V]\|_F. (100) Using ‖X⊙Y‖F≤‖X‖F‖Y‖F\|X Y\|_F≤\|X\|_F\|Y\|_F and ‖X⊙Y‖F≤‖X‖max‖Y‖F\|X Y\|_F≤\|X\|_ \|Y\|_F, we obtain: ‖D2ht(W)[U,V]‖F≤ \|D^2h_t(W)[U,V]\|_F≤ ‖D2A¯t(W)[U,V]‖F‖ht−1(W)‖F \|D^2 A_t(W)[U,V]\|_F\|h_t-1(W)\|_F +‖DA¯t(W)[U]‖F‖Dht−1(W)[V]‖F +\|D A_t(W)[U]\|_F\|Dh_t-1(W)[V]\|_F +‖DA¯t(W)[V]‖F‖Dht−1(W)[U]‖F +\|D A_t(W)[V]\|_F\|Dh_t-1(W)[U]\|_F +‖A¯t(W)‖max‖D2ht−1(W)[U,V]‖F +\| A_t(W)\|_ \|D^2h_t-1(W)[U,V]\|_F +‖D2bt(W)[U,V]‖F. +\|D^2b_t(W)[U,V]\|_F. (101) Then, we already derived the bounds: ‖ht−1(W)‖F≤HR‖Dht−1(W)[U]‖F≤JR‖Dht−1(W)[V]‖F≤JR, \|h_t-1(W) \|_F≤ H_R \|Dh_t-1(W)[U]\|_F≤ J_R \|Dh_t-1(W)[V]\|_F≤ J_R, and by Definition 4.1: ‖A¯t(W)‖max≤ρ, \| A_t(W) \|_ ≤ρ, so the bound becomes: ‖D2ht(W)[U,V]‖F≤ \|D^2h_t(W)[U,V]\|_F≤ HR‖D2A¯t(W)[U,V]‖F H_R\|D^2 A_t(W)[U,V]\|_F +JR‖DA¯t(W)[U]‖F+JR‖DA¯t(W)[V]‖F +J_R\|D A_t(W)[U]\|_F+J_R\|D A_t(W)[V]\|_F +ρ‖D2ht−1(W)[U,V]‖F+‖D2bt(W)[U,V]‖F. +ρ\|D^2h_t-1(W)[U,V]\|_F+\|D^2b_t(W)[U,V]\|_F. (102) Next, we derive the three missing bounds in (102): ‖DA¯t(W)[U]‖F,‖D2A¯t(W)[U,V]‖F,‖D2bt(W)[U,V]‖F \|D A_t(W)[U] \|_F, \|D^2 A_t(W)[U,V] \|_F, \|D^2b_t(W)[U,V] \|_F (103) noting that the same upper bound applies to ‖DA¯t(W)[U]‖F\|D A_t(W)[U]\|_F and ‖DA¯t(W)[V]‖F\|D A_t(W)[V]\|_F. We first derive bounds on the intermediate terms that comprise A¯t A_t, and B¯t B_t, starting with Δt(W) _t(W). From the definition of Δt(W) _t(W): DΔt(W)[UΔ]=diag⁡(Softplus′⁡(WΔxt))UΔxt,D _t(W)[U_ ]=diag (Softplus (W_ x_t) )U_ x_t, (104) and since 0≤Softplus′⁡(z)≤10 (z)≤ 1, it follows that: ‖DΔt(W)[UΔ]‖2≤‖UΔxt‖2≤‖UΔ‖F‖xt‖2≤Lx. \|D _t(W)[U_ ] \|_2≤ \|U_ x_t \|_2≤\|U_ \|_F\|x_t\|_2≤ L_x. (105) Similarly, D2Δt(W)[UΔ,VΔ]=Softplus′⁡(WΔxt)⊙(UΔxt)⊙(VΔxt).D^2 _t(W)[U_ ,V_ ]=Softplus (W_ x_t) (U_ x_t) (V_ x_t). (106) Since 0≤Softplus′⁡(z)≤1/40 (z)≤ 1/4, the bound becomes: ‖D2Δt(W)[UΔ,VΔ]‖2≤14‖(UΔxt)⊙(VΔxt)‖2≤14‖UΔxt‖2‖VΔxt‖2≤14Lx2. \|D^2 _t(W)[U_ ,V_ ] \|_2≤ 14 \|(U_ x_t) (V_ x_t) \|_2≤ 14\|U_ x_t\|_2\|V_ x_t\|_2≤ 14L_x^2. These bounds allow us to bound DA¯t(W)[U]D A_t(W)[U]. We define an intermediate matrix Zt(W):=AΔt(W)⊤,Z_t(W):=A _t(W) , (107) and calculate its derivative DZt(W)[U]=UAΔt(W)⊤+A(DΔt(W)[UΔ])⊤.DZ_t(W)[U]=U_A _t(W) +A (D _t(W)[U_ ] ) . (108) Next, using the property of Frobenius norm for outer product ‖xy⊤‖F=‖x‖2‖y‖2\|xy \|_F=\|x\|_2\|y\|_2, it follows that: ‖DZt(W)[U]‖F≤‖UA‖2‖Δt(W)‖2+‖A‖2‖DΔt(W)[UΔ]‖2≤MKΔ+MALx. \|DZ_t(W)[U] \|_F≤\|U_A\|_2\| _t(W)\|_2+\|A\|_2 \|D _t(W)[U_ ] \|_2≤ M_K_ +M_AL_x. (109) This allows us to calculate: ‖DA¯t(W)[U]‖F \|D A_t(W)[U] \|_F =‖A¯t(W)⊙DZt(W)[U]‖F = \| A_t(W) DZ_t(W)[U] \|_F ≤‖A¯t(W)‖max‖DZt(W)[U]‖F≤ρ⁡(MKΔ+MALx):=KA1,R ≤ \| A_t(W) \|_ \|DZ_t(W)[U] \|_F≤ρ (M_K_ +M_AL_x ):=K_A_1,R (110) Next, we bound ‖D2A¯t(W)[U,V]‖F \|D^2 A_t(W)[U,V] \|_F by differentiating DA¯t(W)[U]D A_t(W)[U] in direction V: D2A¯t(W)[U,V]=A¯t(W)⊙(DZt(W)[V]⊙DZt(W)[U]+D2Zt(W)[U,V]).D^2 A_t(W)[U,V]= A_t(W) (DZ_t(W)[V] DZ_t(W)[U]+D^2Z_t(W)[U,V] ). (111) To do that we first bound: D2Zt(W)[U,V]=UA(DΔt(W)[VΔ])⊤+VA(DΔt(W)[UΔ])⊤+A(D2Δt(W)[UΔ,VΔ])⊤.D^2Z_t(W)[U,V]=U_A (D _t(W)[V_ ] ) +V_A (D _t(W)[U_ ] ) +A (D^2 _t(W)[U_ ,V_ ] ) . (112) We then apply previously derived bounds to the individual components: ‖D2Zt(W)[U,V]‖F≤ \|D^2Z_t(W)[U,V] \|_F≤ ‖UA‖2‖DΔt(W)[VΔ]‖2 \|U_A\|_2 \|D _t(W)[V_ ] \|_2 +‖VA‖2‖DΔt(W)[UΔ]‖2+‖A‖2‖D2Δt(W)[UΔ,VΔ]‖2 +\|V_A\|_2 \|D _t(W)[U_ ] \|_2+\|A\|_2 \|D^2 _t(W)[U_ ,V_ ] \|_2 ≤ ≤ 2Lx+14MALx2. 2L_x+ 14M_AL_x^2. (113) Plugging the result back in, we get the second component bound: ‖D2A¯t(W)[U,V]‖F \|D^2 A_t(W)[U,V] \|_F ≤‖A¯t(W)‖max(‖DZt(W)[V]‖F‖DZt(W)[U]‖F+‖D2Zt(W)[U,V]‖F) ≤ \| A_t(W) \|_ ( \|DZ_t(W)[V] \|_F \|DZ_t(W)[U] \|_F+ \|D^2Z_t(W)[U,V] \|_F ) ≤ρ⁡((MKΔ+MALx)2+2Lx+14MALx2):=KA2,R ≤ρ ( (M_K_ +M_AL_x )^2+2L_x+ 14M_AL_x^2 ):=K_A_2,R (114) Finally, we bound the third term D2bt(W)D^2b_t(W). From definition, bt(W)=B¯t(W)diag⁡(xt)b_t(W)= B_t(W)diag(x_t), where B¯t(W)=(WBxt)Δt(W)⊤ B_t(W)=(W_Bx_t) _t(W) . We first compute the second derivative of B¯t(W) B_t(W): D2B¯t(W)[U,V] D^2 B_t(W)[U,V] =(UBxt)(DΔt(W)[VΔ])⊤+(VBxt)(DΔt(W)[UΔ])⊤+(WBxt)(D2Δt(W)[UΔ,VΔ])⊤, =(U_Bx_t) (D _t(W)[V_ ] ) +(V_Bx_t)(D _t(W)[U_ ]) +(W_Bx_t)(D^2 _t(W)[U_ ,V_ ]) , (115) and using the previously derived bounds, it follows that: ‖D2B¯t(W)[U,V]‖F≤2Lx2+14MBLx3. \|D^2 B_t(W)[U,V] \|_F≤ 2L_x^2+ 14M_BL_x^3. (116) With that, we get: ‖D2bt(W)[U,V]‖F \|D^2b_t(W)[U,V] \|_F ≤‖D2B¯t(W)[U,V]‖F‖diag⁡(xt)‖F ≤ \|D^2 B_t(W)[U,V] \|_F \|diag(x_t) \|_F ≤(2Lx2+14MBLx3)Lx=2Lx3+14MBLx4:=KB2,R ≤ (2L_x^2+ 14M_BL_x^3 )L_x=2L_x^3+ 14M_BL_x^4:=K_B_2,R (117) Given these derived bounds, the recurrence becomes: ‖D2ht(W)[U,V]‖F≤ρ‖D2ht−1(W)[U,V]‖F+KA2,RHR+2KA1,RJR+KB2,R. \|D^2h_t(W)[U,V] \|_F≤ρ \|D^2h_t-1(W)[U,V] \|_F+K_A_2,RH_R+2K_A_1,RJ_R+K_B_2,R. (118) Since, as mentioned in the proof of 4.4, the initial state is initialized as h0=0h_0=0, we have: D2h0(W)[U,V]=0.D^2h_0(W)[U,V]=0. (119) As such, unrolling the recurrence we get: ‖D2ht(W)[U,V]‖F \|D^2h_t(W)[U,V] \|_F ≤(KA2,RHR+2KA1,RJR+KB2,R)∑m=0t−1ρm ≤ (K_A_2,RH_R+2K_A_1,RJ_R+K_B_2,R ) _m=0^t-1ρ^m ≤KA2,RHR+2KA1,RJR+KB2,R1−ρ:=KH2,R ≤ K_A_2,RH_R+2K_A_1,RJ_R+K_B_2,R1-ρ:=K_H_2,R (120) Therefore, for all W∈R(Wr)W _R(W^r) and all unit block perturbations U,VU,V, ‖D2ht(W)[U,V]‖F≤KH2,R. \|D^2h_t(W)[U,V] \|_F≤ K_H_2,R. (121) Next, we use this to derive the first and second Fréchet derivatives of the output: yt(W)=ht(W)⊤(WCxt).y_t(W)=h_t(W) (W_Cx_t). (122) The first derivative is: Dyt(W)[U]=Dht(W)[U]⊤Ct(W)+ht(W)⊤UCxt,Dy_t(W)[U]=Dh_t(W)[U] C_t(W)+h_t(W) U_Cx_t, (123) resulting in a bound ‖Dyt(W)[U]‖2 \|Dy_t(W)[U] \|_2 ≤‖Dht(W)[U]‖F‖WC‖2‖xt‖2+‖ht(W)‖F‖UCxt‖2 ≤ \|Dh_t(W)[U] \|_F \|W_C \|_2\|x_t\|_2+\|h_t(W)\|_F\|U_Cx_t\|_2 ≤JRMCLx+HRLx:=KY1,R. ≤ J_RM_CL_x+H_RL_x:=K_Y_1,R. (124) Taking the Fréchet derivative in direction V, we get: D2yt(W)[U,V]=D2ht(W)[U,V]⊤Ct(W)+Dht(W)[U]⊤VCxt+Dht(W)[V]⊤UCxt, D^2y_t(W)[U,V]=D^2h_t(W)[U,V] C_t(W)+Dh_t(W)[U] V_Cx_t+Dh_t(W)[V] U_Cx_t, (125) where we also use the fact that D2(WCxt)=0D^2(W_Cx_t)=0. Then: ‖D2yt(W)[U,V]‖2≤ \|D^2y_t(W)[U,V] \|_2≤ ‖D2ht(W)[U,V]‖F‖(WCxt)‖2 \|D^2h_t(W)[U,V]\|_F\|(W_Cx_t)\|_2 +‖Dht(W)[U]‖F‖(VCxt)‖2+‖Dht(W)[V]‖F‖(UCxt)‖2 +\|Dh_t(W)[U]\|_F\|(V_Cx_t)\|_2+\|Dh_t(W)[V]\|_F\|(U_Cx_t)\|_2 ≤ ≤ KH2,RMCLx+2JRLx:=KY2,R. K_H_2,RM_CL_x+2J_RL_x:=K_Y_2,R. (126) Finally, consider the token loss ℓt(yt(W)) _t(y_t(W)). By the second-order chain rule: DW2ℓt(yt(W))[U,V]=Dy2ℓt(yt)[Dyt(W)[U],Dyt(W)[V]]+Dyℓt(yt)[D2yt(W)[U,V]], D^2_W _t(y_t(W))[U,V]=D^2_y _t(y_t)[Dy_t(W)[U],Dy_t(W)[V]]+D_y _t(y_t)[D^2y_t(W)[U,V]], (127) where we use the notation Dy2ℓt(yt)[a,b]=a⊤∇y2ℓt(yt)bD^2_y _t(y_t)[a,b]=a _y^2 _t(y_t)b. By Assumption 4.2, and utilizing all derived bounds, we get: |DW2ℓt(yt(W))[U,V]| |D_W^2 _t(y_t(W))[U,V] | ≤|Dy2ℓt(yt)[Dyt(W)[U],Dyt(W)[V]]|+|Dyℓt(yt)[D2yt(W)[U,V]]| ≤ |D_y^2 _t(y_t)[Dy_t(W)[U],Dy_t(W)[V]] |+ |D_y _t(y_t)[D^2y_t(W)[U,V]] | ≤‖∇y2ℓt(yt)‖op‖Dyt(W)[U]‖2‖Dyt(W)[V]‖2+‖∇yℓt(yt)‖2‖D2yt(W)[U,V]‖2 ≤\| _y^2 _t(y_t)\|_op\|Dy_t(W)[U]\|_2\|Dy_t(W)[V]\|_2+\| _y _t(y_t)\|_2\|D^2y_t(W)[U,V]\|_2 ≤HℓKY1,R2+LℓKY2,R:=LSSM,R. ≤ H_ K_Y_1,R^2+L_ K_Y_2,R:=L_SSM,R. (128) Substituting the definitions, we get: LSSM,R= L_SSM,R= HℓLx2(HR+MCJR)2+LℓLx(MCKH2,R+2JR) H_ L_x^2 (H_R+M_CJ_R )^2+L_ L_x (M_CK_H_2,R+2J_R ) = = HℓLx2(MKBMKΔND+MCMKW1−ρ)2 H_ L_x^2 ( M_K_BM_K_ ND+M_CM_K_W1-ρ )^2 +LℓLx[MC1−ρ(ρ((MKΔ+MALx)2+2Lx+14MALx2)MKBMKΔND1−ρ +L_ L_x [ M_C1-ρ (ρ ((M_K_ +M_AL_x)^2+2L_x+ 14M_AL_x^2 ) M_K_BM_K_ ND1-ρ +2ρ(MKΔ+MALx)MKW1−ρ+2Lx3+14MBLx4)+2MKW1−ρ]. +2ρ(M_K_ +M_AL_x) M_K_W1-ρ+2L_x^3+ 14M_BL_x^4 )+ 2M_K_W1-ρ ]. (129) Since U,VU,V were arbitrary unit block perturbations, the previous bound implies ‖∇W2ℓt(yt(W))‖op≤LSSM,R. \| _W^2 _t(y_t(W)) \|_op≤ L_SSM,R. (130) Finally, the sequence loss was defined in (3) as ℒS(W)=1S∑t=1Sℓt(yt(W)),L^S(W)= 1S _t=1^S _t(y_t(W)), using the simplified notation, where ℒS(W):=ℒS(W,ξ)L^S(W):=L^S(W;ξ) and ℓt(yt(W))=ℓt(yt(W,x),yt∗) _t(y_t(W))= _t(y_t(W;x),y_t^*), and we further incorporate W to emphasize yty_t’s dependence on the current weights. As a result, ‖∇W2ℒS(W)‖op \| _W^2L^S(W)\|_op =‖1S∑t=1S∇W2ℓt(yt(W))‖op = \| 1S _t=1^S _W^2 _t(y_t(W)) \|_op ≤1S∑t=1S‖∇W2ℓt(yt(W))‖op ≤ 1S _t=1^S\| _W^2 _t(y_t(W))\|_op ≤LSSM,R. ≤ L_SSM,R. (131) Then, for an client k, the local expected objective is ℒk(W)=ξ∼k[ℒS(W,ξ)].L_k(W)=E_ξ _k [L^S(W;ξ) ]. (132) For any W,V∈R(Wr)W,V _R(W^r), convexity of the neighborhood implies that Wα:=V+α⁡(W−V),α∈[0,1],W_α:=V+α(W-V), α∈[0,1], (133) lies in R(Wr)N_R(W^r). Applying the fundamental theorem of calculus to the fixed-sequence gradient and using the Hessian bound from (131) gives ‖∇WℒS(W,ξ)−∇WℒS(V,ξ)‖F \| _WL^S(W;ξ)- _WL^S(V;ξ) \|_F =‖vec⁡(∇WℒS(W,ξ)−∇WℒS(V,ξ))‖2 = \|vec ( _WL^S(W;ξ)- _WL^S(V;ξ) ) \|_2 =‖∫01∇W2ℒS(Wα,ξ)vec⁡(W−V)α‖2 = \| _0^1 _W^2L^S(W_α;ξ)vec(W-V)dα \|_2 ≤(a)∫01‖∇W2ℒS(Wα,ξ)vec⁡(W−V)‖2α (a)≤ _0^1 \| _W^2L^S(W_α;ξ)vec(W-V) \|_2dα ≤(b)∫01‖∇W2ℒS(Wα,ξ)‖op‖W−V‖Fα (b)≤ _0^1 \| _W^2L^S(W_α;ξ) \|_op \|W-V \|_Fdα ≤(c)∫01LSSM,R‖W−V‖Fα=LSSM,R‖W−V‖F. (c)≤ _0^1L_SSM,R\|W-V\|_F\,dα=L_SSM,R\|W-V\|_F. (134) where: • (a)(a) follows from the triangle inequality for an integral ‖∫v‖≤∫‖v‖ \| v \|≤ \|v\|. • (b)(b) follows from the operator-norm inequality: ‖Hd‖2≤‖H‖op‖d‖2\|Hd\|_2≤\|H\|_op\|d\|_2 and the fact that ‖vec⁡(W−V)‖2=‖W−V‖F \|vec(W-V) \|_2= \|W-V \|_F. • (c)(c) follows from (131). This inequality holds uniformly over all admissible sequence pairs and finite sequence lengths. Finally, by Assumption 4.6 and Jensen’s inequality, ‖∇Wℒk(W)−∇Wℒk(V)‖F \| _WL_k(W)- _WL_k(V)\|_F =‖ξ∼k[∇WℒS(W,ξ)−∇WℒS(V,ξ)]‖F = \|E_ξ _k [ _WL^S(W;ξ)- _WL^S(V;ξ) ] \|_F ≤ξ∼k[‖∇WℒS(W,ξ)−∇WℒS(V,ξ)‖F] _ξ _k [\| _WL^S(W;ξ)- _WL^S(V;ξ)\|_F ] ≤LSSM,R‖W−V‖F. ≤ L_SSM,R\|W-V\|_F. (135) Therefore, ℒkL_k is LSSM,RL_SSM,R-smooth on R(Wr)N_R(W^r). ∎ B.1.3 SSM loss gradient second-moment constant Lemma and Proof In the following two lemmas, we consider some arbitrary client data distribution P and write ℒ⁡(W)=ξ∼[ℒS(W,ξ)]L(W)=E_ξ [L^S(W;ξ) ]. Expectations written as ξE_ξ are taken with respect to the distribution P. Lemma B.2 (SSM loss gradient second-moment constant). Suppose the Assumptions 4.2 and 4.3 hold. Let R(Wr)⊆γ,δN_R(W^r) _γ,δ be a stable local model neighborhood around the round-r broadcast model WrW^r. Let: g⁡(W,ξ):=∇WℒS(W,ξ)g(W;ξ):= _WL^S(W;ξ) (136) denote the stochastic sequence-level gradient for a sampled input-target sequence pair ξ=(x1:S,y1:S∗)ξ=(x_1:S,y_1:S^*). Then, for every W∈R(Wr)W _R(W^r), ξ‖g⁡(W,ξ)‖F2≤GSSM,R2,E_ξ \|g(W;ξ) \|_F^2≤ G_SSM,R^2, (137) where GSSM,R:=LℓLx(MKBMKΔND1−ρ+MCMKW1−ρ),G_SSM,R:=L_ L_x (M_K_BM_K_ ND1-ρ+M_C M_K_W1-ρ ), (138) where MKB,MKΔ,MC, and MKWM_K_B,M_K_ ,M_C, and M_K_W are defined in (66), (65), (59), and (84) respectively. Proof of Lemma B.2: By Theorem 4.4, under the assumptions of the lemma, for every W∈R(Wr)W _R(W^r), every finite integer S≥1S≥ 1, and every admissible length-S sequence pair ξ∈ξ , we have: ‖∇WℒS(W,ξ)‖F≤LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ). \| _WL^S(W;ξ) \|_F≤ L_ L_x (K_BK_ ND1-ρ+\|W_C\|_2 K_W1-ρ ). Since g⁡(W,ξ)=∇WℒS(W,ξ)g(W;ξ)= _WL^S(W;ξ), it follows that: ‖g⁡(W,ξ)‖F≤LℓLx(KBKΔND1−ρ+‖WC‖2KW1−ρ).\|g(W;ξ)\|_F≤ L_ L_x (K_BK_ ND1-ρ+\|W_C\|_2 K_W1-ρ ). (139) Restricting to W∈R(Wr)W _R(W^r), by the local neighborhood bounds defined previously ( (66), (65), (59), and (84)), KB≤MKB,KΔ≤MKΔ,KW≤MKW,‖WC‖2≤MC. K_B≤ M_K_B, K_ ≤ M_K_ , K_W≤ M_K_W, \|W_C\|_2≤ M_C. (140) Therefore, ‖g⁡(W,ξ)‖F≤LℓLx(MKBMKΔND1−ρ+MCMKW1−ρ)=GSSM,R.\|g(W;ξ)\|_F≤ L_ L_x (M_K_BM_K_ ND1-ρ+M_C M_K_W1-ρ )=G_SSM,R. (141) Since both values are non-negative, squaring both sides gives: ‖g⁡(W,ξ)‖F2≤GSSM,R2=Lℓ2Lx2(MKBMKΔND1−ρ+MCMKW1−ρ)2\|g(W;ξ)\|_F^2≤ G_SSM,R^2=L_ ^2L_x^2 (M_K_BM_K_ ND1-ρ+M_C M_K_W1-ρ )^2 (142) Finally, taking expectation over the sampled sequence ξ preserves the bound: ξ[‖g⁡(W,ξ)‖F2]≤GSSM,R2E_ξ [\|g(W;ξ)\|_F^2 ]≤ G_SSM,R^2 (143) proving the lemma. ∎ B.1.4 SSM stochastic gradient variance bound Lemma and Proof Lemma B.3 (SSM stochastic gradient variance bound). Suppose the assumptions of Lemma B.2 and Assumption 4.6 hold. Let ℒ⁡(W):=ξ[ℒS(W,ξ)]L(W):=E_ξ [L^S(W;ξ) ] denote the expected sequence loss over the local data distribution, and let g⁡(W,ξ):=∇WℒS(W,ξ)g(W;ξ):= _WL^S(W;ξ) denote the single-sequence stochastic gradient. For a mini-batch Ξξ=ξ1,…,ξΞ _ξ=\ _1,…, _ \ of Ξ independent sampled sequences, define gΞ(W,Ξξ):=1Ξ∑j=1Ξg⁡(W,ξj).g_ (W; _ξ):= 1 _j=1 g(W; _j). (144) Then, for every W∈R(Wr)W _R(W^r) it follows that Ξξ‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2≤σSSM,R,Ξ2,E_ _ξ \|g_ (W; _ξ)- _WL(W) \|_F^2≤ _SSM,R, ^2, (145) where σSSM,R,Ξ2:=1ΞsupW∈R(Wr)(ξ‖g⁡(W,ξ)‖F2−‖∇Wℒ(W)‖F2). _SSM,R, ^2:= 1 _W _R(W^r) (E_ξ\|g(W;ξ)\|_F^2-\| _WL(W)\|_F^2 ). (146) Moreover, by Lemma B.2, σSSM,R,Ξ2≤GSSM,R2Ξ. _SSM,R, ^2≤ G_SSM,R^2 . (147) Proof of Lemma B.3: By unbiasedness and independence of the sampled sequences, the standard mini-batch variance identity gives: Ξξ[‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2]=1Ξξ[‖g⁡(W,ξ)−∇Wℒ(W)‖F2]. _ _ξ [ \|g_ (W; _ξ)- _WL(W) \|_F^2 ]= 1 E_ξ [ \|g(W;ξ)- _WL(W) \|_F^2 ]. Using the identity ξ[‖g⁡(W,ξ)−∇Wℒ(W)‖F2]=ξ‖g⁡(W,ξ)‖F2−‖∇Wℒ(W)‖F2,E_ξ [ \|g(W;ξ)- _WL(W) \|_F^2 ]=E_ξ\|g(W;ξ)\|_F^2-\| _WL(W)\|_F^2, (148) and taking the supremum over W∈R(Wr)W _R(W^r), we obtain Ξξ[‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2]≤σSSM,R,Ξ2.E_ _ξ [ \|g_ (W; _ξ)- _WL(W) \|_F^2 ]≤ _SSM,R, ^2. (149) Finally, Lemma B.2 gives ξ‖g⁡(W,ξ)‖F2≤GSSM,R2,E_ξ\|g(W;ξ)\|_F^2≤ G_SSM,R^2, and since ‖∇Wℒ(W)‖F2≥0\| _WL(W)\|_F^2≥ 0, σSSM,R,Ξ2≤GSSM,R2Ξ. _SSM,R, ^2≤ G_SSM,R^2 . (150) This proves the claim. ∎ B.1.5 SSM client heterogeneity bound under total variation Lemma and Proof Lemma B.4 (SSM client heterogeneity bound under total variation). Suppose the assumptions of Lemma B.2 and Assumption 4.6 hold. Let f⁡(W):=∑k=1Kpkℒk(W)f(W):= _k=1^Kp_kL_k(W) (151) be the global objective, where pk≥0p_k≥ 0 and ∑k=1Kpk=1 _k=1^Kp_k=1. Let kP_k denote the local data distribution of client k, and define the global mixture distribution as: ¯:=∑k=1Kpkk. P:= _k=1^Kp_kP_k. (152) Define the local client heterogeneity constant ΓSSM,R2:=supW∈R(Wr)∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2. _SSM,R^2:= _W _R(W^r) _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2. (153) Then: ΓSSM,R2≤4GSSM,R2∑k=1KpkTV⁡(k,¯)2, _SSM,R^2≤ 4G_SSM,R^2 _k=1^Kp_kTV(P_k, P)^2, (154) where TV⁡(P,Q):=supA|P⁡(A)−Q⁡(A)|.TV(P,Q):= _A|P(A)-Q(A)|. (155) Proof of Lemma B.4: Recall that for each client k, the local objective is: ℒk(W)=ξ∼k[ℒS(W,ξ)],L_k(W)=E_ξ _k [L^S(W;ξ) ], and, by Assumption 4.6, ∇Wℒk(W)=ξ∼k[g⁡(W,ξ)]. _WL_k(W)=E_ξ _k [g(W;ξ) ]. (156) Similarly, since f⁡(W):=∑k=1Kpkℒk(W),f(W):= _k=1^Kp_kL_k(W), we have ∇Wf(W)=∑k=1Kpk∇Wℒk(W)=∑k=1Kpkξ∼k[g⁡(W,ξ)]. _Wf(W)= _k=1^Kp_k _WL_k(W)= _k=1^Kp_kE_ξ _k [g(W;ξ) ]. (157) Next, by the definition of the mixture distribution ¯ P, we can write this as: ∇Wf(W)=ξ∼¯[g⁡(W,ξ)]. _Wf(W)=E_ξ P [g(W;ξ) ]. (158) Therefore, it follows that ∇Wℒk(W)−∇Wf(W)=ξ∼k[g⁡(W,ξ)]−ξ∼¯[g⁡(W,ξ)]. _WL_k(W)- _Wf(W)=E_ξ _k[g(W;ξ)]-E_ξ P[g(W;ξ)]. (159) By the pointwise bound (141) established in the proof of Lemma B.2, for every W∈R(Wr)W _R(W^r) and every admissible sequence pair ξ∈ξ , ‖g⁡(W,ξ)‖F≤GSSM,R. \|g(W;ξ) \|_F≤ G_SSM,R. Since k()=¯()=1P_k(Z)= P(Z)=1, the total-variation inequality for bounded vector-valued functions gives: ‖ξ∼k[g⁡(W,ξ)]−ξ∼¯[g⁡(W,ξ)]‖F≤2GSSM,RTV⁡(k,¯). \|E_ξ _k[g(W;ξ)]-E_ξ P[g(W;ξ)] \|_F≤ 2G_SSM,RTV(P_k, P). (160) So: ‖∇Wℒk(W)−∇Wf(W)‖F2≤4GSSM,R2TV⁡(k,¯)2. \| _WL_k(W)- _Wf(W) \|_F^2≤ 4G_SSM,R^2TV(P_k, P)^2. (161) Multiplying both sides by pkp_k and summing over all clients gives: ∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2≤4GSSM,R2∑k=1KpkTV⁡(k,¯)2. _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2≤ 4G_SSM,R^2 _k=1^Kp_kTV(P_k, P)^2. (162) Since the right-hand side of this inequality is independent of W, taking the supremum over W∈R(Wr)W _R(W^r) we get: ΓSSM,R2:=supW∈R(Wr)∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2≤4GSSM,R2∑k=1KpkTV⁡(k,¯)2. _SSM,R^2:= _W _R(W^r) _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2≤ 4G_SSM,R^2 _k=1^Kp_kTV(P_k, P)^2. (163) This proves the SSM client heterogeneity bound. ∎ B.2 Theorem B.5: FedAvg SSM Convergence and Proof Theorem B.5 (FedAvg SSM Convergence). Consider a distributed learning framework established in Section 3, with the objective f⁡(W):=∑k=1Kpkℒk(W),pk≥0,∑k=1Kpk=1,f(W):= _k=1^Kp_kL_k(W), p_k≥ 0, _k=1^Kp_k=1, where W parameterizes a single-layer selective SSM defined in the same section and each ℒkL_k is finite and differentiable on the considered neighborhoods. Let W0W^0 be deterministic with f⁡(W0)<∞f(W^0)<∞ and suppose f is bounded below by a finite constant f∗f^*. Consider full-participation FedAvg distributed training algorithm applied to f⁡(W)f(W) for T≥1T≥ 1 communication rounds, E≥1E≥ 1 local steps per round, Ξ≥1 ≥ 1 mini-batch size, and constant deterministic local step size η>0η>0, with the global model evolution described as: Wkr,0 W_k^r,0 =Wr, =W^r, Wkr,e+1 W_k^r,e+1 =Wkr,e−ηgk,Ξr,e for e∈0,…,E−1, =W_k^r,e-η g_k, ^r,e~~ for ~~e∈\0,…,E-1\, Wr+1 W^r+1 =∑k=1KpkWkr,E = _k=1^Kp_kW_k^r,E where gk,Ξr,eg_k, ^r,e is the average of Ξ sequence-level gradients evaluated at Wkr,eW_k^r,e with sampling and conditional unbiasedness, and E,Ξ,T∈ℤ>0E, ,T _>0. Suppose this model evolution satisfies, almost surely: Wkr,e∈R(Wr)⊆γ,δ,‖WXr‖2+R≤MX∀X∈Δ,B,C,A,W_k^r,e _R(W^r) _γ,δ,~~~~ \|W_X^r \|_2+R≤ M_X~~∀~~X∈\ ,B,C,A\, for every client k and local step 0≤e≤E0≤ e≤ E, where WA:=AW_A:=A, and R,γ,δ>0R,γ,δ>0 and the finite deterministic envelopes MXM_X are common across rounds. Let Assumptions 4.2- 4.6 hold uniformly on these neighborhoods. Define ρ:=exp⁡(−γδ)ρ:= (-γδ) and auxiliary SSM-specific quantities as follows: MKΔ M_K_ :=MΔLx+Dlog⁡2, :=M_ L_x+ D 2, MKB M_K_B :=MBLx2, :=M_BL_x^2, MKW M_K_W :=NDMKBMKΔ(MKΔ+DMALx)ρ1−ρ+Lx2NMKΔ+LxDMKB, :=NDM_K_BM_K_ (M_K_ + DM_AL_x ) ρ1-ρ+L_x^2 NM_K_ +L_x DM_K_B, HR H_R :=MKBMKΔND1−ρ,JR:=MKW1−ρ, := M_K_BM_K_ ND1-ρ, J_R:= M_K_W1-ρ, MA1,R M_A_1,R :=ρ⁡(MKΔ+MALx), :=ρ(M_K_ +M_AL_x), MA2,R M_A_2,R :=ρ⁡[(MKΔ+MALx)2+2Lx+14MALx2], :=ρ [(M_K_ +M_AL_x)^2+2L_x+ 14M_AL_x^2 ], MB2,R M_B_2,R :=2Lx3+14MBLx4, :=2L_x^3+ 14M_BL_x^4, MH2,R M_H_2,R :=MA2,RHR+2MA1,RJR+MB2,R1−ρ. := M_A_2,RH_R+2M_A_1,RJ_R+M_B_2,R1-ρ. These quantities determine the common SSM-specific smoothness and gradient bounds: LSSM,R L_SSM,R :=HℓLx2(HR+MCJR)2+LℓLx(MCMH2,R+2JR), :=H_ L_x^2(H_R+M_CJ_R)^2+L_ L_x(M_CM_H_2,R+2J_R), GSSM,R G_SSM,R :=LℓLx(HR+MCJR). :=L_ L_x(H_R+M_CJ_R). For the client distributions, further define ¯:=∑k=1Kpkk,ℋTV:=∑k=1KpkTV⁡(k,¯)2, P:= _k=1^Kp_kP_k, _TV:= _k=1^Kp_kTV (P_k, P )^2, where TV⁡(P,Q):=sup|P⁡()−Q⁡()|TV(P,Q):= _A |P(A)-Q(A) | is the total variation distance, with the supremum taken over measurable sets A. Then, for any constant deterministic step size η>0η>0 satisfying ηELSSM,R≤12,12η2LSSM,R2E(E−1)≤1,η EL_SSM,R≤ 12, 12η^2L_SSM,R^2E(E-1)≤ 1, the following holds: 1T∑r=0T−1‖∇f(Wr)‖F2≤ 1T _r=0^T-1E\|∇ f(W^r)\|_F^2≤ 4(f⁡(W0)−f∗)ηET+4ηLSSM,RGSSM,R2Ξ∑k=1Kpk2 4 (f(W^0)-f^* )η ET+ 4η L_SSM,RG_SSM,R^2 _k=1^Kp_k^2 +6η2LSSM,R2(E−1)GSSM,R2Ξ+48η2LSSM,R2E(E−1)GSSM,R2ℋTV. + 6η^2L_SSM,R^2(E-1)G_SSM,R^2 +48η^2L_SSM,R^2E(E-1)G_SSM,R^2H_TV. Proof of Theorem B.5: We start from a standard nonconvex FedAvg convergence bound applicable to the setting considered here. Previous works establish a general bound (see Supplemental Material C.6, Equation (96) of Wang et al. [2020]) for smooth local objectives, unbiased stochastic gradients with bounded variance, and bounded client gradient dissimilarity. For a setting with full client participation, E local SGD steps per client, and objective and aggregation weights pkp_k, the bound in Wang et al. [2020] takes the following form: 1T∑r=0T−1‖∇f(Wr)‖F2≤4(f⁡(W0)−f∗)ηET+4ηLσ2∑k=1Kpk2+6η2L2(E−1)σ2+12η2L2E(E−1)Γ2, 1T _r=0^T-1E\|∇ f(W^r)\|_F^2≤ 4 (f(W^0)-f^* )η ET+4η Lσ^2 _k=1^Kp_k^2+6η^2L^2(E-1)σ^2+12η^2L^2E(E-1) ^2, (164) whenever η>0η>0 satisfies ηLE≤12,12η2L2E(E−1)≤1.η LE≤ 12, 12η^2L^2E(E-1)≤ 1. Here, L is a common smoothness constant, σ2σ^2 bounds the conditional variance of each mini-batch gradient estimator, and Γ2 ^2 bounds the weighted client gradient disagreement: ∑k=1Kpk‖∇ℒk(W)−∇f(W)‖F2≤Γ2. _k=1^Kp_k\| _k(W)-∇ f(W)\|_F^2≤ ^2. (165) Note that the inequality (164) specializes the result in Wang et al. [2020] applied to full client participation, E local SGD steps at every client, and aggregation weights equal to the objective weights pkp_k. We apply this bound to the selective SSM setting by substituting the constants established in our lemmas. First, Lemma 4.7, evaluated with the common envelopes MXM_X, gives L=LSSM,R,L=L_SSM,R, (166) with LSSM,RL_SSM,R defined in Theorem B.5. Since each neighborhood R(Wr)N_R(W^r) is convex and contains all local iterates in round r, it also contains their weighted averages and the segment joining WrW^r to Wr+1W^r+1. As such, all smoothness inequalities used in the convergence argument hold with the common constant LSSM,RL_SSM,R. Next, Lemma B.2 derives the common stochastic gradient second-moment bound GSSM,R2G_SSM,R^2. Together with the sampling conditions in Assumption 4.6, Lemma B.3 establishes the choice σ2=GSSM,R2Ξ.σ^2= G_SSM,R^2 . (167) Finally, Lemma B.4 establishes the common heterogeneity bound Γ2=4GSSM,R2∑k=1KpkTV⁡(k,¯)2=4GSSM,R2ℋTV. ^2=4G_SSM,R^2 _k=1^Kp_kTV(P_k, P)^2=4G_SSM,R^2H_TV. (168) Substituting these quantities into (164), we obtain the inequality stated in Theorem B.5 whenever η>0η>0 satisfies the stated conditions with L=LSSM,RL=L_SSM,R. ∎ B.3 Theorem B.6: FedProx SSM Convergence and Proof The FedProx algorithm modifies the local client objective by adding a proximal penalty around the round-r broadcast model WrW^r: Fkr(W):=ℒk(W)+μ2‖W−Wr‖F2,μ>0,F_k^r(W):=L_k(W)+ μ2 \|W-W^r \|_F^2, μ>0, (169) thus changing the local optimization dynamics. In this section we introduce a theorem that quantifies the convergence behavior of a distributed learning framework that utilizes the FedProx with selective State Space Models, and provide a corresponding proof. Theorem B.6 (FedProx convergence for selective SSMs). Consider a distributed learning framework established in Section 3, with the objective f⁡(W):=∑k=1Kpkℒk(W),pk≥0,∑k=1Kpk=1,f(W):= _k=1^Kp_kL_k(W), p_k≥ 0, _k=1^Kp_k=1, where W parameterizes a single-layer selective SSM defined in the same section and each ℒkL_k is finite and differentiable on the considered neighborhoods. Let W0W^0 be deterministic with f⁡(W0)<∞f(W^0)<∞ and suppose f is bounded below by a finite constant f∗f^*. Consider full-participation FedProx distributed training algorithm applied to f⁡(W)f(W) for T≥1T≥ 1 communication rounds, E≥1E≥ 1 local steps per round, Ξ≥1 ≥ 1 mini-batch size, constant deterministic local step size η>0η>0, and constant deterministic proximal coefficient μ>0μ>0, with the global model evolution described as: Wkr,0 W_k^r,0 =Wr, =W^r, Wkr,e+1 W_k^r,e+1 =Wkr,e−η⁡(gk,Ξr,e+μ⁡(Wkr,e−Wr)) for e∈0,…,E−1, =W_k^r,e-η (g_k, ^r,e+μ (W_k^r,e-W^r ) )~~ for ~~e∈\0,…,E-1\, Wr+1 W^r+1 =∑k=1KpkWkr,E = _k=1^Kp_kW_k^r,E where gk,Ξr,eg_k, ^r,e is the average of Ξ sequence-level gradients evaluated at Wkr,eW_k^r,e with sampling and conditional unbiasedness, and E,Ξ,T∈ℤ>0E, ,T _>0. Suppose that, almost surely for every round r<Tr<T: Wkr,e∈R(Wr)⊆γ,δ,‖WXr‖2+R≤MX∀X∈Δ,B,C,A,W_k^r,e _R(W^r) _γ,δ,~~~~ \|W_X^r \|_2+R≤ M_X~~∀~~X∈\ ,B,C,A\, (170) for every client k and local step 0≤e≤E0≤ e≤ E, where WA:=AW_A:=A, and R,γ,δ>0R,γ,δ>0 and the finite deterministic envelopes MXM_X are common across rounds. Let Assumptions 4.2- 4.6 hold uniformly on these neighborhoods. Define ρ:=exp⁡(−γδ)ρ:= (-γδ). Assume 0<ημ≤10<ημ≤ 1, and define the effective radius RμR_μ by Rμ:=min⁡R,ηGSSM,R∑j=0E−1(1−ημ)j=min⁡R,GSSM,Rμ[1−(1−ημ)E],R_μ:= \R,η G_SSM,R _j=0^E-1(1-ημ)^j \= \R, G_SSM,Rμ [1-(1-ημ)^E ] \, where GSSM,RG_SSM,R is defined in Theorem B.5, with the corresponding FedProx parameter envelopes defined as: MXμ:=MX−R+Rμ∀X∈Δ,B,C,A.M_X^μ:=M_X-R+R_μ~~∀~~X∈\ ,B,C,A\. The corresponding auxiliary SSM-specific quantities are then defined as: MKΔμ M_K_ ^μ :=MΔμLx+Dlog⁡2, :=M_ ^μL_x+ D 2, MKBμ M_K_B^μ :=MBμLx2, :=M_B^μL_x^2, MKWμ M_K_W^μ :=NDMKBμMKΔμ(MKΔμ+DMAμLx)ρ1−ρ+Lx2NMKΔμ+LxDMKBμ, :=NDM_K_B^μM_K_ ^μ (M_K_ ^μ+ DM_A^μL_x ) ρ1-ρ+L_x^2 NM_K_ ^μ+L_x DM_K_B^μ, HRμ H_R_μ :=MKBμMKΔμND1−ρ,JRμ:=MKWμ1−ρ, := M_K_B^μM_K_ ^μ ND1-ρ, J_R_μ:= M_K_W^μ1-ρ, MA1,Rμ M_A_1,R^μ :=ρ⁡(MKΔμ+MAμLx), :=ρ(M_K_ ^μ+M_A^μL_x), MA2,Rμ M_A_2,R^μ :=ρ⁡[(MKΔμ+MAμLx)2+2Lx+14MAμLx2], :=ρ [(M_K_ ^μ+M_A^μL_x)^2+2L_x+ 14M_A^μL_x^2 ], MB2,Rμ M_B_2,R^μ :=2Lx3+14MBμLx4, :=2L_x^3+ 14M_B^μL_x^4, MH2,Rμ M_H_2,R^μ :=MA2,RμHRμ+2MA1,RμJRμ+MB2,Rμ1−ρ. := M_A_2,R^μH_R_μ+2M_A_1,R^μJ_R_μ+M_B_2,R^μ1-ρ. Let: LSSM,Rμ L_SSM,R_μ :=HℓLx2(HRμ+MCμJRμ)2+LℓLx(MCμMH2,Rμ+2JRμ), :=H_ L_x^2(H_R_μ+M_C^μJ_R_μ)^2+L_ L_x(M_C^μM^μ_H_2,R+2J_R_μ), Lprox,Rμ L_prox,R_μ :=LSSM,Rμ+μ. :=L_SSM,R_μ+μ. GSSM,Rμ G_SSM,R_μ :=LℓLx(HRμ+MCμJRμ), :=L_ L_x(H_R_μ+M_C^μJ_R_μ), denote common deterministic SSM-specific smoothness and gradient bounds constants evaluated on Rμ(Wr)N_R_μ(W^r). For the client distributions, further define ¯:=∑k=1Kpkk,ℋTV:=∑k=1KpkTV⁡(k,¯)2, P:= _k=1^Kp_kP_k, _TV:= _k=1^Kp_kTV (P_k, P )^2, where TV⁡(P,Q):=sup|P⁡()−Q⁡()|TV(P,Q):= _A |P(A)-Q(A) | is the total variation distance, with the supremum taken over measurable sets A. Then, for any constant deterministic step size η>0η>0 satisfying ηELSSM,Rμ≤12,12η2Lprox,Rμ2E(E−1)≤1,η EL_SSM,R_μ≤ 12, 12η^2L_prox,R_μ^2E(E-1)≤ 1, the following bound holds: 1T∑r=0T−1‖∇f(Wr)‖F2≤ 1T _r=0^T-1E \|∇ f(W^r) \|_F^2≤ 3(f⁡(W0)−f∗)ηET+3ηLSSM,RμGSSM,Rμ2Ξ∑k=1Kpk2 3 (f(W^0)-f^* )η ET+ 3η L_SSM,R_μG_SSM,R_μ^2 _k=1^Kp_k^2 +3η2Lprox,Rμ2(E−1)GSSM,Rμ2Ξ+24η2Lprox,Rμ2E(E−1)GSSM,Rμ2ℋTV. + 3η^2L_prox,R_μ^2(E-1)G_SSM,R_μ^2 +24η^2L_prox,R_μ^2E(E-1)G_SSM,R_μ^2H_TV. Moreover, almost surely, for every round r<Tr<T, client k, and local step 0≤e≤E0≤ e≤ E, ‖Wkr,e−Wr‖F≤min⁡R,GSSM,Rμ[1−(1−ημ)e]≤Rμ. \|W_k^r,e-W^r \|_F≤ \R, G_SSM,Rμ [1-(1-ημ)^e ] \≤ R_μ. Consequently, all local iterates and the aggregated model Wr+1W^r+1 lie in Rμ(Wr)⊆γ,δN_R_μ(W^r) _γ,δ. Proof of Theorem B.6: The FedProx framework modifies the local update as: Wkr,e+1=Wkr,e−η⁡(gk,Ξr,e+μ⁡(Wkr,e−Wr)),Wkr,0=Wr.W_k^r,e+1=W_k^r,e-η (g_k, ^r,e+μ(W_k^r,e-W^r) ), W_k^r,0=W^r. (171) Subtracting the broadcast model WrW^r gives Wkr,e+1−Wr=(1−ημ)(Wkr,e−Wr)−ηgk,Ξr,eW_k^r,e+1-W^r=(1-ημ)(W_k^r,e-W^r)-η g_k, ^r,e (172) Taking Frobenius norms, using 0<ημ≤10<ημ≤ 1, and by the triangle inequality, ‖Wkr,e+1−Wr‖F≤(1−ημ)‖Wkr,e−Wr‖F+η‖gk,Ξr,e‖F. \|W_k^r,e+1-W^r \|_F≤(1-ημ) \|W_k^r,e-W^r \|_F+η \|g_k, ^r,e \|_F. (173) By the bound (141) established in the proof of Lemma B.2, every admissible sequence gradient satisfies ‖gk(W,ξ)‖F≤GSSM,R \|g_k(W;ξ) \|_F≤ G_SSM,R for all W∈R(Wr)W _R(W^r). As such, the mini-batch gradients satisfy: ‖gk,Ξr,e‖F=‖1Ξ∑i=1Ξgk(Wkr,e,ξk,ir,e)‖F≤1Ξ∑i=1Ξ‖gk(Wkr,e,ξk,ir,e)‖F≤GSSM,R \|g_k, ^r,e \|_F= \| 1 _i=1 g_k(W^r,e_k;ξ^r,e_k,i) \|_F≤ 1 _i=1 \|g_k(W^r,e_k;ξ^r,e_k,i) \|_F≤ G_SSM,R (174) Therefore: ‖Wkr,e+1−Wr‖F≤(1−ημ)‖Wkr,e−Wr‖F+ηGSSM,R. \|W_k^r,e+1-W^r \|_F≤(1-ημ) \|W_k^r,e-W^r \|_F+η G_SSM,R. (175) Unrolling the recursion from Wkr,0=WrW_k^r,0=W^r for every 0≤e≤E0≤ e≤ E, we get: ‖Wkr,e−Wr‖F≤ηGSSM,R∑e′=0e−1(1−ημ)e′=GSSM,Rμ[1−(1−ημ)e] \|W_k^r,e-W^r \|_F≤η G_SSM,R _e =0^e-1(1-ημ)^e = G_SSM,Rμ [1-(1-ημ)^e ] (176) Combining this with the established neighborhood bounds ‖Wkr,e−Wr‖F≤R \|W_k^r,e-W^r \|_F≤ R, we get: ‖Wkr,e−Wr‖F≤min⁡(R,GSSM,Rμ[1−(1−ημ)e])≤Rμ, \|W_k^r,e-W^r \|_F≤ (R, G_SSM,Rμ [1-(1-ημ)^e ] )≤ R_μ, (177) and since pk≥0p_k≥ 0 and ∑kpk=1 _kp_k=1, it follows that: ‖Wr+1−Wr‖F=‖∑k=1Kpk(Wkr,E−Wr)‖F≤∑k=1Kpk‖Wkr,E−Wr‖F≤Rμ. \|W^r+1-W^r \|_F= \| _k=1^Kp_k (W_k^r,E-W^r ) \|_F≤ _k=1^Kp_k \|W_k^r,E-W^r \|_F≤ R_μ. (178) This shows that all local iterates as well as Wr+1W^r+1 belong to Rμ(Wr)⊆γ,δN_R_μ(W^r) _γ,δ. Since this neighborhood is convex, the segments used below in the smoothness arguments also remain inside it. Next, for every W∈Rμ(Wr)W _R_μ(W^r) and parameter block X∈Δ,B,C,AX∈\ ,B,C,A\, ‖WX‖2≤‖WXr‖2+Rμ≤MX−R+Rμ=:MXμ. \|W_X \|_2≤\|W_X^r\|_2+R_μ≤ M_X-R+R_μ=:M_X^μ. (179) These MXμM_X^μ are uniform across rounds and deterministic. Let LSSM,RμL_SSM,R_μ and GSSM,RμG_SSM,R_μ denote the SSM smoothness and stochastic-gradient constants constructed using MXμM_X^μ, with the stability, input, and loss-regularity constants remaining unchanged. In particular: MKΔμ M_K_ ^μ :=MΔμLx+Dlog⁡2, :=M_ ^μL_x+ D 2, (180) MKBμ M_K_B^μ :=MBμLx2, :=M_B^μL_x^2, (181) MKWμ M_K_W^μ :=NDMKBμMKΔμ(MKΔμ+DMAμLx)ρ1−ρ+Lx2NMKΔμ+LxDMKBμ, :=NDM_K_B^μM_K_ ^μ (M_K_ ^μ+ DM_A^μL_x ) ρ1-ρ+L_x^2 NM_K_ ^μ+L_x DM_K_B^μ, (182) HRμ H_R_μ :=MKBμMKΔμND1−ρ,JRμ:=MKWμ1−ρ, := M_K_B^μM_K_ ^μ ND1-ρ, J_R_μ:= M_K_W^μ1-ρ, (183) MA1,Rμ M_A_1,R^μ :=ρ⁡(MKΔμ+MAμLx), :=ρ(M_K_ ^μ+M_A^μL_x), (184) MA2,Rμ M_A_2,R^μ :=ρ⁡[(MKΔμ+MAμLx)2+2Lx+14MAμLx2], :=ρ [(M_K_ ^μ+M_A^μL_x)^2+2L_x+ 14M_A^μL_x^2 ], (185) MB2,Rμ M_B_2,R^μ :=2Lx3+14MBμLx4, :=2L_x^3+ 14M_B^μL_x^4, (186) MH2,Rμ M_H_2,R^μ :=MA2,RμHRμ+2MA1,RμJRμ+MB2,Rμ1−ρ, := M_A_2,R^μH_R_μ+2M_A_1,R^μJ_R_μ+M_B_2,R^μ1-ρ, (187) LSSM,Rμ L_SSM,R_μ :=HℓLx2(HRμ+MCμJRμ)2+LℓLx(MCμMH2,Rμ+2JRμ), :=H_ L_x^2(H_R_μ+M_C^μJ_R_μ)^2+L_ L_x(M_C^μM_H_2,R^μ+2J_R_μ), (188) GSSM,Rμ G_SSM,R_μ :=LℓLx(HRμ+MCμJRμ), :=L_ L_x(H_R_μ+M_C^μJ_R_μ), (189) and in addition define: Lprox,Rμ=LSSM,Rμ+μL_prox,R_μ=L_SSM,R_μ+μ (190) On this neighborhood, Lemma 4.7 implies that each ℒkL_k is LSSM,RμL_SSM,R_μ-smooth. Then, for any W,V∈Rμ(Wr)W,V _R_μ(W^r), the gradients of the local proximal objective satisfy: ‖(∇ℒk(W)+μ⁡(W−Wr))−(∇ℒk(V)+μ⁡(V−Wr))‖F \| ( _k(W)+μ(W-W^r) )- ( _k(V)+μ(V-W^r) ) \|_F ≤‖∇ℒk(W)−∇ℒk(V)‖F+μ‖W−V‖F ≤\| _k(W)- _k(V)\|_F+μ\|W-V\|_F ≤(LSSM,Rμ+μ)‖W−V‖F ≤ (L_SSM,R_μ+μ ) \|W-V \|_F =Lprox,Rμ‖W−V‖F. =L_prox,R_μ\|W-V\|_F. (191) Therefore, the local proximal objective ℒk(W)+μ2‖W−Wr‖F2L_k(W)+ μ2\|W-W^r\|_F^2 is Lprox,RμL_prox,R_μ-smooth on Rμ(Wr)N_R_μ(W^r). Moreover, since f=∑kpkℒkf= _kp_kL_k with pk≥0p_k≥ 0 and ∑kpk=1 _kp_k=1, the global objective f remains LSSM,RμL_SSM,R_μ-smooth on the same neighborhood. Next, define the following common deterministic variance constant: σSSM,Rμ2:=GSSM,Rμ2Ξ.σ^2_SSM,R_μ:= G_SSM,R_μ^2 . (192) By Lemma B.4, for every round r<Tr<T, supW∈Rμ(Wr)∑kpk‖∇ℒk(W)−∇f(W)‖F2≤4GSSM,Rμ2ℋTV. _W _R_μ(W^r) _kp_k\| _k(W)-∇ f(W)\|_F^2≤ 4G_SSM,R_μ^2H_TV. (193) The right-hand side provides a common deterministic heterogeneity bound across rounds. Next, fix a round r<Tr<T and let r[⋅]E_r[·] denote an expectation conditional on the training history at the beginning of that round, with WrW^r being fixed under this conditioning. We begin by unrolling the local update and subtracting the starting weights: Wkr,e−Wr W_k^r,e-W^r =−η∑j=0e−1[gk,Ξr,j+μ(Wkr,j−Wr)] =-η _j=0^e-1 [g^r,j_k, +μ (W^r,j_k-W^r ) ] =−η∑j=0e−1[gk,Ξr,j+μ(Wkr,j−Wr)] =-η _j=0^e-1 [g^r,j_k, +μ (W^r,j_k-W^r ) ] +η∑j=0e−1(∇ℒk(Wr)−∇ℒk(Wkr,j))−η∑j=0e−1(∇ℒk(Wr)−∇ℒk(Wkr,j)) +η _j=0^e-1 ( _k (W^r )- _k (W^r,j_k ) )-η _j=0^e-1 ( _k (W^r )- _k (W^r,j_k ) ) =−ηe∇ℒk(Wr)−η∑j=0e−1[gk,Ξr,j−∇ℒk(Wkr,j)] =-η e _k (W^r )-η _j=0^e-1 [g^r,j_k, - _k (W^r,j_k ) ] −η∑j=0e−1[μ(Wkr,j−Wr)−∇ℒk(Wr)+∇ℒk(Wkr,j)]. -η _j=0^e-1 [μ (W^r,j_k-W^r )- _k (W^r )+ _k (W^r,j_k ) ]. (194) Next, when evaluating the norms of the individual quantities on the right of (194), by Lemma B.3, the following holds r[‖∑j=0e−1(gk,Ξr,j−∇ℒk(Wkr,j))‖F2]≤eσSSM,Rμ2E_r [ \| _j=0^e-1 (g_k, ^r,j- _k (W_k^r,j ) ) \|_F^2 ]≤ e _SSM,R_μ^2 (195) for all 0≤e≤E0≤ e≤ E. Then, by substituting Wkr,jW_k^r,j and WrW^r into (191), it follows that: ‖μ⁡(Wkr,j−Wr)−∇ℒk(Wr)+∇ℒk(Wkr,j)‖F≤Lprox,Rμ‖Wkr,j−Wr‖F \|μ (W^r,j_k-W^r )- _k (W^r )+ _k (W^r,j_k ) \|_F≤ L_prox,R_μ \|W^r,j_k-W^r \|_F (196) Taking the norm of (194), applying the squared-normed inequality for a sum of three terms as well as the Cauchy-Schwarz inequality, and substituting these results, we get: r[‖Wkr,e−Wr‖F2] _r [ \|W_k^r,e-W^r \|_F^2 ] ≤3η2e2‖∇ℒk(Wr)‖F2+3η2eσSSM,Rμ2 ≤ 3η^2e^2 \| _k (W^r ) \|_F^2+3η^2e _SSM,R_μ^2 +3η2eLprox,Rμ2∑j=0e−1r[‖Wkr,j−Wr‖F2]. +3η^2eL_prox,R_μ^2 _j=0^e-1E_r [ \|W^r,j_k-W^r \|_F^2 ]. (197) By multiplying (197) by pkp_k and summing over all clients, we get: ∑k=1Kpkr[‖Wkr,e−Wr‖F2] _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] ≤3η2eσSSM,Rμ2+3η2e2∑k=1Kpk‖∇ℒk(Wr)‖F2 ≤ 3η^2e _SSM,R_μ^2+3η^2e^2 _k=1^Kp_k \| _k (W^r ) \|_F^2 +3η2eLprox,Rμ2∑j=0e−1∑k=1Kpkr[‖Wkr,j−Wr‖F2]. +3η^2eL_prox,R_μ^2 _j=0^e-1 _k=1^Kp_kE_r [ \|W^r,j_k-W^r \|_F^2 ]. (198) Next, note that: ∑k=1Kpk‖∇ℒk(Wr)‖F2 _k=1^Kp_k \| _k (W^r ) \|_F^2 =∑k=1Kpk‖∇ℒk(Wr)+∇f(Wr)−∇f(Wr)‖F2 = _k=1^Kp_k \| _k (W^r )+∇ f(W^r)-∇ f(W^r) \|_F^2 ≤‖∇f(Wr)‖F2+∑k=1Kpk‖∇ℒk(Wr)−∇f(Wr)‖F2 ≤ \|∇ f(W^r) \|_F^2+ _k=1^Kp_k \| _k (W^r )-∇ f(W^r) \|_F^2 ≤‖∇f(Wr)‖F2+4GSSM,Rμ2ℋTV. ≤ \|∇ f(W^r) \|_F^2+4G_SSM,R_μ^2H_TV. (199) where the last inequality follows from (193). Substituting this result into (198), we get: ∑k=1Kpkr[‖Wkr,e−Wr‖F2] _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] ≤3η2eσSSM,Rμ2+3η2e2(‖∇f(Wr)‖F2+4GSSM,Rμ2ℋTV) ≤ 3η^2e _SSM,R_μ^2+3η^2e^2 ( \|∇ f(W^r) \|_F^2+4G_SSM,R_μ^2H_TV ) +3η2eLprox,Rμ2∑j=0e−1∑k=1Kpkr[‖Wkr,j−Wr‖F2]. +3η^2eL_prox,R_μ^2 _j=0^e-1 _k=1^Kp_kE_r [ \|W^r,j_k-W^r \|_F^2 ]. (200) Averaging (200) over e=0,…,E−1e=0,…,E-1, where we utilize 1E∑e=0E−1e=E−12 1E _e=0^E-1e= E-12 and 1E∑e=0E−1e2≤E⁡(E−1) 1E _e=0^E-1e^2≤ E(E-1), yields: 1E∑e=1E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] 1E _e=1^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] ≤32η2(E−1)σSSM,Rμ2 ≤ 32η^2(E-1) _SSM,R_μ^2 +3η2E(E−1)(‖∇f(Wr)‖F2+4GSSM,Rμ2ℋTV) +3η^2E(E-1) ( \|∇ f(W^r) \|_F^2+4G_SSM,R_μ^2H_TV ) +3η2E(E−1)Lprox,Rμ2[1E∑e=0E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2]], +3η^2E(E-1)L_prox,R_μ^2 [ 1E _e=0^E-1 _k=1^Kp_kE_r [ \|W^r,e_k-W^r \|_F^2 ] ], (201) where for the last term, we used: 1E∑e=0E−1e∑j=0e−1∑k=1Kpkr[‖Wkr,j−Wr‖F2]≤E⁡(E−1)[1E∑e=0E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2]] 1E _e=0^E-1e _j=0^e-1 _k=1^Kp_kE_r [ \|W^r,j_k-W^r \|_F^2 ]≤ E(E-1) [ 1E _e=0^E-1 _k=1^Kp_kE_r [ \|W^r,e_k-W^r \|_F^2 ] ] (202) Rearranging the terms in (201) and dividing out by 1−3η2E(E−1)Lprox,Rμ21-3η^2E(E-1)L_prox,R_μ^2, yields: 1E∑e=1E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] 1E _e=1^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] ≤11−3η2E(E−1)Lprox,Rμ2(32η2(E−1)σSSM,Rμ2CLOSE ≤ 11-3η^2E(E-1)L_prox,R_μ^2 ( 32η^2(E-1) _SSM,R_μ^2 OPEN+3η2E(E−1)(‖∇f(Wr)‖F2+4GSSM,Rμ2ℋTV)). +3η^2E(E-1) ( \|∇ f(W^r) \|_F^2+4G_SSM,R_μ^2H_TV ) ). (203) Next, note that by the step-size condition 12η2E(E−1)Lprox,Rμ2≤112η^2E(E-1)L_prox,R_μ^2≤ 1, it is true that 3η2E(E−1)Lprox,Rμ2≤143η^2E(E-1)L_prox,R_μ^2≤ 14, and therefore 11−3η2E(E−1)Lprox,Rμ2≤43 11-3η^2E(E-1)L_prox,R_μ^2≤ 43. Substituting that result into (203) yields: 1E∑e=1E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] 1E _e=1^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] ≤2η2(E−1)σSSM,Rμ2 ≤ 2η^2(E-1) _SSM,R_μ^2 +4η2E(E−1)(‖∇f(Wr)‖F2+4GSSM,Rμ2ℋTV). +4η^2E(E-1) ( \|∇ f(W^r) \|_F^2+4G_SSM,R_μ^2H_TV ). (204) Next, we evaluate Wr+1−WrW^r+1-W^r. By the update rule: Wr+1−Wr=−η∑e=0E−1∑k=1Kpk[gk,Ξr,e+μ(Wkr,e−Wr)].W^r+1-W^r=-η _e=0^E-1 _k=1^Kp_k [g^r,e_k, +μ (W^r,e_k-W^r ) ]. (205) Since the segment joining WrW^r and Wr+1W^r+1 lies in Rμ(Wr)N_R_μ(W^r), by the smoothness of f established in Lemma 4.7, it is true that: r[f⁡(Wr+1)]≤ _r [f(W^r+1) ]≤ f⁡(Wr) f(W^r) −η⟨∇f(Wr),r[∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))]⟩F -η ∇ f (W^r ),E_r [ _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r ) ) ] _F +η2LSSM,Rμ2r[‖∑e=0E−1∑k=1Kpk(gk,Ξr,e+μ⁡(Wkr,e−Wr))‖F2]. + η^2L_SSM,R_μ2E_r [ \| _e=0^E-1 _k=1^Kp_k (g^r,e_k, +μ (W^r,e_k-W^r ) ) \|_F^2 ]. (206) Here, the tower property and because of conditional unbiasedness, the stochastic gradients can be replaced by local gradients inside the expected inner product. Next, note that: ∑e=0E−1∑k=1Kpk(gk,Ξr,e+μ⁡(Wkr,e−Wr)) _e=0^E-1 _k=1^Kp_k (g^r,e_k, +μ (W^r,e_k-W^r ) ) =∑e=0E−1∑k=1Kpk(gk,Ξr,e+μ⁡(Wkr,e−Wr)+∇ℒk(Wkr,e)−∇ℒk(Wkr,e)) = _e=0^E-1 _k=1^Kp_k (g^r,e_k, +μ (W^r,e_k-W^r )+ _k (W_k^r,e )- _k (W_k^r,e ) ) =∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr)) = _e=0^E-1 _k=1^Kp_k ( _k (W_k^r,e )+μ (W^r,e_k-W^r ) ) +∑e=0E−1∑k=1Kpk(gk,Ξr,e−∇ℒk(Wkr,e)). + _e=0^E-1 _k=1^Kp_k (g^r,e_k, - _k (W_k^r,e ) ). (207) When taking norms, for a sum of two matrices, the squared norm is at most twice the sum of their squared norms. As such, it follows that: r[‖∑e=0E−1∑k=1Kpk(gk,Ξr,e+μ⁡(Wkr,e−Wr))‖F2] _r [ \| _e=0^E-1 _k=1^Kp_k (g^r,e_k, +μ (W^r,e_k-W^r ) ) \|_F^2 ] ≤2r[‖∑e=0E−1∑k=1Kpk(gk,Ξr,e−∇ℒk(Wkr,e))‖F2] ≤ 2E_r [ \| _e=0^E-1 _k=1^Kp_k (g^r,e_k, - _k (W_k^r,e ) ) \|_F^2 ] +2r[‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))‖F2] +2E_r [ \| _e=0^E-1 _k=1^Kp_k ( _k (W_k^r,e )+μ (W^r,e_k-W^r ) ) \|_F^2 ] (208) Next, the centered gradient-noise differences have conditional mean zero given the complete history before each local step. Conditional independence across clients eliminates the same-step cross terms, while conditional unbiasedness and the tower property eliminate the cross terms between different local steps. Applying Lemma B.3 therefore gives r[‖∑e=0E−1∑k=1Kpk(gk,Ξr,e−∇ℒk(Wkr,e))‖F2]≤∑e=0E−1∑k=1Kpk2σSSM,Rμ2=EσSSM,Rμ2∑k=1Kpk2,E_r [ \| _e=0^E-1 _k=1^Kp_k (g^r,e_k, - _k (W_k^r,e ) ) \|_F^2 ]≤ _e=0^E-1 _k=1^Kp_k^2 _SSM,R_μ^2=E _SSM,R_μ^2 _k=1^Kp_k^2, (209) and as a result: r[‖∑e=0E−1∑k=1Kpk(gk,Ξr,e+μ⁡(Wkr,e−Wr))‖F2] _r [ \| _e=0^E-1 _k=1^Kp_k (g^r,e_k, +μ (W^r,e_k-W^r ) ) \|_F^2 ] ≤2EσSSM,Rμ2∑k=1Kpk2 ≤ 2E _SSM,R_μ^2 _k=1^Kp_k^2 +2r[‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))‖F2] +2E_r [ \| _e=0^E-1 _k=1^Kp_k ( _k (W_k^r,e )+μ (W^r,e_k-W^r ) ) \|_F^2 ] (210) When it comes to the inner-product term, we note that, since f⁡(Wr)=∑k=1Kpkℒk(Wr)f(W^r)= _k=1^Kp_kL_k(W^r), it is true that: E∇f(Wr)=∑e=0E−1∑k=1Kpk∇ℒk(Wr).E∇ f(W^r)= _e=0^E-1 _k=1^Kp_k _k(W^r). (211) Using this result, and the −2⟨U,V⟩F=‖U−V‖F2−‖U‖F2−‖V‖F2-2 U,V _F=\|U-V\|_F^2-\|U\|_F^2-\|V\|_F^2 equality, the following holds: − - η⟨∇f(Wr),r[∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))]⟩F η ∇ f (W^r ),E_r [ _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r ) ) ] _F (212) =−ηE2‖∇f(Wr)‖F2 =- η E2 \|∇ f(W^r) \|_F^2 −η2Er[‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))‖F2] - η2EE_r [ \| _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r ) ) \|_F^2 ] +η2Er[‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr)−∇ℒk(Wr))‖F2], + η2EE_r [ \| _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r )- _k (W^r ) ) \|_F^2 ], (213) where the E-terms appeared from evaluating the inner-product with E∇f(Wr)E∇ f(W^r) and dividing out by E. Then, the positive term is bounded as: r _r [‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr)−∇ℒk(Wr))‖F2] [ \| _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r )- _k (W^r ) ) \|_F^2 ] ≤E∑e=0E−1∑k=1Kpkr[‖∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr)−∇ℒk(Wr)‖F2] ≤ E _e=0^E-1 _k=1^Kp_kE_r [ \| _k (W^r,e_k )+μ (W_k^r,e-W^r )- _k (W^r ) \|_F^2 ] ≤ELprox,Rμ2∑e=0E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] ≤ EL_prox,R_μ^2 _e=0^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] (214) where the first inequality is true by Jensen’s inequality with weights pk/Ep_k/E, and the second inequality follows from (191). Substituting these results back into (206), we obtain: r[f⁡(Wr+1)]≤ _r [f (W^r+1 ) ]≤ f⁡(Wr)−ηE2‖∇f(Wr)‖F2 f (W^r )- η E2 \|∇ f(W^r) \|_F^2 +(η2LSSM,Rμ−η2E)r[‖∑e=0E−1∑k=1Kpk(∇ℒk(Wkr,e)+μ⁡(Wkr,e−Wr))‖F2] + (η^2L_SSM,R_μ- η2E )E_r [ \| _e=0^E-1 _k=1^Kp_k ( _k (W^r,e_k )+μ (W_k^r,e-W^r ) ) \|_F^2 ] +ηLprox,Rμ22∑e=0E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] + η L_prox,R_μ^22 _e=0^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] +η2ELSSM,RμσSSM,Rμ2∑k=1Kpk2. +η^2EL_SSM,R_μ _SSM,R_μ^2 _k=1^Kp_k^2. (215) The step-size condition ηELSSM,Rμ≤1/2η EL_SSM,R_μ≤ 1/2 implies that η2LSSM,Rμ−η2E≤0η^2L_SSM,R_μ- η2E≤ 0, and therefore the term multiplied by this factor is never positive. As such, the following holds: r[f⁡(Wr+1)]≤ _r [f (W^r+1 ) ]≤ f⁡(Wr)−ηE2‖∇f(Wr)‖F2 f (W^r )- η E2 \|∇ f(W^r) \|_F^2 +ηLprox,Rμ22∑e=0E−1∑k=1Kpkr[‖Wkr,e−Wr‖F2] + η L_prox,R_μ^22 _e=0^E-1 _k=1^Kp_kE_r [ \|W_k^r,e-W^r \|_F^2 ] +η2ELSSM,RμσSSM,Rμ2∑k=1Kpk2. +η^2EL_SSM,R_μ _SSM,R_μ^2 _k=1^Kp_k^2. (216) Substituting the result from (204) into this inequality, we get: r[f⁡(Wr+1)]≤ _r [f (W^r+1 ) ]≤ f⁡(Wr) f (W^r ) +(2η3E2(E−1)Lprox,Rμ2−ηE2)‖∇f(Wr)‖F2 + (2η^3E^2(E-1)L_prox,R_μ^2- η E2 ) \|∇ f(W^r) \|_F^2 +η3E(E−1)Lprox,Rμ2σSSM,Rμ2 +η^3E(E-1)L_prox,R_μ^2 _SSM,R_μ^2 +8η3E2(E−1)Lprox,Rμ2GSSM,Rμ2ℋTV +8η^3E^2(E-1)L_prox,R_μ^2G_SSM,R_μ^2H_TV +η2ELSSM,RμσSSM,Rμ2∑k=1Kpk2. +η^2EL_SSM,R_μ _SSM,R_μ^2 _k=1^Kp_k^2. (217) Next, by the second step size condition 12η2Lprox,Rμ2E(E−1)≤112η^2L_prox,R_μ^2E(E-1)≤ 1 it follows that: 2η3E2(E−1)Lprox,Rμ2−ηE2 2η^3E^2(E-1)L_prox,R_μ^2- η E2 =ηE(2η2E(E−1)Lprox,Rμ2−12) =η E (2η^2E(E-1)L_prox,R_μ^2- 12 ) ≤ηE(16−12)=−ηE3. ≤η E ( 16- 12 )=- η E3. (218) Substituting in (217), we get: r[f⁡(Wr+1)]≤ _r [f (W^r+1 ) ]≤ f⁡(Wr) f (W^r ) −ηE3‖∇f(Wr)‖F2 - η E3 \|∇ f(W^r) \|_F^2 +η3E(E−1)Lprox,Rμ2σSSM,Rμ2 +η^3E(E-1)L_prox,R_μ^2 _SSM,R_μ^2 +8η3E2(E−1)Lprox,Rμ2GSSM,Rμ2ℋTV +8η^3E^2(E-1)L_prox,R_μ^2G_SSM,R_μ^2H_TV +η2ELSSM,RμσSSM,Rμ2∑k=1Kpk2. +η^2EL_SSM,R_μ _SSM,R_μ^2 _k=1^Kp_k^2. (219) Taking total expectations in (219), rearranging, and summing over r=0,…,T−1r=0,…,T-1 makes the objective differences telescope. Using ⁡[f⁡(WT)]≥f∗E[f(W^T)]≥ f^* and dividing by ηET/3η ET/3, we obtain: 1T∑r=0T−1‖∇f(Wr)‖F2≤ 1T _r=0^T-1E \|∇ f(W^r) \|_F^2≤ 3(f⁡(W0)−f∗)ηET+3ηLSSM,RμσSSM,Rμ2∑k=1Kpk2 3 (f(W^0)-f^* )η ET+3η L_SSM,R_μ _SSM,R_μ^2 _k=1^Kp_k^2 +3η2(E−1)Lprox,Rμ2σSSM,Rμ2+24η2E(E−1)Lprox,Rμ2GSSM,Rμ2ℋTV. +3η^2(E-1)L_prox,R_μ^2 _SSM,R_μ^2+24η^2E(E-1)L_prox,R_μ^2G_SSM,R_μ^2H_TV. (220) Finally, substituting σSSM,Rμ2=GSSM,Rμ2/Ξ _SSM,R_μ^2=G_SSM,R_μ^2/ (which follows from Lemma B.3) into (220), gives the stated bound, completing the proof. ∎ Appendix C Multi-Layer Selective SSM proofs In this section we derive the corresponding Gradient, FedAvg, and FedProx convergence bounds for a Selective SSM with of an arbitrary finite depth Q, extending our derivations in Sections A, and B to a multi-layer selective State Space Model (SSM) case. C.1 Multi-Layer Selective SSM Gradient Bound Proof Similar to the single-layer setting, we first derive an upper bound on the gradient norm of the loss function for an arbitrary sequence of length S processed by a selective state space model (SSM) composed of Q sequential selective SSM layers. The analysis once again utilizes definitions from Section 3 and an input sequence x=(x1,…,xS)x=(x_1,…,x_S) where xt∈ℝD×1x_t ^D× 1. The multi-layer selective SSM consists of Q layers arranged sequentially. Each layer q∈1,…,Qq∈ \1,…,Q \ is parameterized by the learnable parameter tuple W(q)=(WΔ(q),WB(q),WC(q),A(q)),W^(q)= (W_ ^(q),W_B^(q),W_C^(q),A^(q) ), (221) where WΔ(q)∈ℝD×D,WB(q)∈ℝN×D,WC(q)∈ℝN×D,A(q)∈ℝN×1W_ ^(q) ^D× D,W_B^(q) ^N× D,W_C^(q) ^N× D,A^(q) ^N× 1. Note that we therefore consider a multi-layer SSM where all of its Q layers have the same input/output dimension D and state dimension N. The full set of model parameters is then given by: W=(W(1),…,W(Q)).W= (W^(1),…,W^(Q) ). (222) In this multi-layer formulation, we denote the input to layer q at time step t by xt(q)x_t^(q), where xt(1)=xtx_t^(1)=x_t and xt(q+1)=yt(q)x_t^(q+1)=y_t^(q) (223) with yt(q)y_t^(q) denoting the output of layer q. Given this formulation, each layer q of the multi-layer selective SSM operates according to the following dynamics: Δt(q) _t^(q) =Softplus(WΔ(q)xt(q)), =Softplus (W_ ^(q)x_t^(q) ), (224) A¯t(q) A^(q)_t =exp⁡(A(q)(Δt(q))⊤), = (A^(q) ( _t^(q) ) ), (225) B¯t(q) B_t^(q) =(WB(q)xt(q))(Softplus(WΔ(q)xt(q)))⊤, = (W_B^(q)x_t^(q) ) (Softplus (W_ ^(q)x_t^(q) ) ) , (226) ht(q) h_t^(q) =A¯t(q)⊙ht−1(q)+B¯t(q)diag(xt(q)), = A^(q)_t h_t-1^(q)+ B_t^(q)diag (x_t^(q) ), (227) yt(q) y_t^(q) =(ht(q))⊤WC(q)xt(q)=(ht(q))⊤Ct(q) = (h_t^(q) ) W_C^(q)x_t^(q)= (h_t^(q) ) C_t^(q) (228) where the hidden state of each layer q is initialized as h0(q)=N×Dh_0^(q)=0_N× D, just as in the single-layer setting. Finally, the empirical loss for processing a sequence of length S for such defined multi-layer selective SSM is: ℒS,(Q)(W,ξ)=1S∑t=1Sℓt(yt(Q))=1S∑t=1Sℓt(xt(Q+1)).L^S,(Q)(W;ξ)= 1S _t=1^S _t (y_t^(Q) )= 1S _t=1^S _t (x_t^(Q+1) ). (229) Here, ξ=(x,y∗)ξ=(x,y^*) is an admissible input-target pair of length S≥1S≥ 1, and ℓt(yt(Q)) _t (y_t^(Q) ) abbreviates ℓt(yt(Q)(W,x),yt∗) _t(y_t^(Q)(W;x),y_t^*). For fixed ξ, we write ℒS,(Q)(W)L^S,(Q)(W) in place of ℒS,(Q)(W,ξ)L^S,(Q)(W;ξ). In this multi-layer setting, the full gradient is the block tuple: ∇WℒS,(Q)(W)=(∇W(1)ℒS,(Q)(W),…,∇W(Q)ℒS,(Q)(W)). _WL^S,(Q)(W)= ( _W^(1)L^S,(Q)(W),…, _W^(Q)L^S,(Q)(W) ). (230) Similar to the single-layer proof, in the multi-layer formulation, gradients with respect to the full parameter tuple W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ) are understood as block gradients across layers. As such, for any function f, ‖∇Wf‖F2:=∑q=1Q‖∇W(q)f‖F2, \| _Wf \|_F^2:= _q=1^Q \| _W^(q)f \|_F^2, (231) where, for each layer q, the gradient is further decomposed into parameter blocks as: ‖∇W(q)f‖F2:=‖∇WΔ(q)f‖F2+‖∇WB(q)f‖F2+‖∇WC(q)f‖F2+‖∇A(q)f‖F2. \| _W^(q)f \|_F^2:= \| _W_ ^(q)f \|_F^2+ \| _W_B^(q)f \|_F^2+ \| _W_C^(q)f \|_F^2+ \| _A^(q)f \|_F^2. (232) As in the single-layer proof, parameter derivatives with respect to W(q)W^(q) are understood as flattened block Jacobians, with rows representing the output coordinates and columns corresponding to the scalar parameters of layer q. We use the same vectorization and scalar-parameter-slice conventions for products and transposes. For this multi-layer formulation, we introduce the following new assumption and definition, that extend the single-layer analysis assumptions to the multi-layer setting: Assumption C.1 (Bounded Layer Inputs). For every layer q∈1,…,Qq∈\1,…,Q\ of a multi-layer selective SSM of depth Q parameterized by a fixed set of weights W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ), there exists a constant Lx,q>0L_x,q>0 such that: ‖xt(q)(W;ξ)‖2≤Lx,q ∀t∈1,…,S. \|x_t^(q) (W;ξ ) \|_2≤ L_x,q ∀ t∈\1,…,S\. (233) This holds for every finite integer length S≥1S≥ 1, every admissible input-target pair ξ∈ξ , and every token index t∈1,…,St∈\1,…,S\, with finite constants independent of S, ξ, and t. Definition C.2 (Multi-layer SSM Stability). For each layer q∈1,…,Qq∈ \1,…,Q \, let γq>0 _q>0 and δq>0 _q>0. We define ,(Q)S_ γ, δ^(Q) as follows: ,(Q):= S_ γ, δ^(Q):= \ W|An(q)≤−γq,[Softplus(WΔ(q)xt(q)(W;ξ))]d≥δq W |A_n^(q)≤- _q, [Softplus (W_ ^(q)x_t^(q) (W;ξ ) ) ]_d≥ _q ∀q∈1,…,Q,∀n∈1,…,N,∀d∈1,…,D,∀t∈1,…,S, ∀ q∈\1,…,Q\,∀ n∈\1,…,N\,∀ d∈\1,…,D\,∀ t∈\1,…,S\, finite integer S≥1,and every admissible length−S pair ξ∈, integer S≥ 1,and every admissible length-S pair ξ \, (234) where =(γ1,…,γQ) γ= ( _1,…, _Q ) and =(δ1,…,δQ) δ= ( _1,…, _Q ) We then say that the full parameter tuple W is stable with these margins if W∈,(Q)W∈S_ γ, δ^(Q). For each layer q, we further define: ρq:=exp⁡(−γqδq)<1. _q:= (- _q _q)<1. (235) As a result of Definition C.2, for a layer q, it follows that: ‖A¯t(q)‖max≤ρq. \| A_t^(q) \|_ ≤ _q. (236) Given such setup, we state the following theorem: Theorem C.3 (Multi-Layer SSM Gradient Bound). Let W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ) be a full parameter tuple satisfying Assumptions 4.2 and C.1, with W∈,(Q)W∈S_ γ, δ^(Q), that parametrizes a multi-layer selective SSM model of finite integer depth Q≥1Q≥ 1. Then, for such W, for every finite integer S≥1S≥ 1 and every admissible input-target pair ξ∈ξ of length S, the following bound holds: ‖∇WℒS,(Q)(W,ξ)‖F≤Lℓ(∑q=1QLx,q2(Hq+‖WC(q)‖2Jq)2∏q′=q+1Qκq′2)1/2:=SSM(Q)(W) \| _WL^S,(Q)(W;ξ) \|_F≤ L_ ( _q=1^QL_x,q^2 (H_q+ \|W_C^(q) \|_2J_q )^2 _q =q+1^Q _q ^2 )^1/2:=G_SSM^(Q)(W) (237) with KΔ,q K_ ,q :=‖WΔ(q)‖2Lx,q+Dlog⁡(2),KB,q:=‖WB(q)‖2Lx,q2, := \|W_ ^(q) \|_2L_x,q+ D (2),~~~~K_B,q:= \|W_B^(q) \|_2L_x,q^2, (238) KW,q K_W,q :=NDKB,qKΔ,q(KΔ,q+D‖A(q)‖2Lx,q)(ρq1−ρq)+Lx,q2NKΔ,q+Lx,qDKB,q, :=NDK_B,qK_ ,q (K_ ,q+ D \|A^(q) \|_2L_x,q ) ( _q1- _q )+L_x,q^2 NK_ ,q+L_x,q DK_B,q, (239) Hq H_q :=KB,qKΔ,qND1−ρq,Jq:=KW,q1−ρq,κq:=‖WC(q)‖2(Hq+Lx,qΛx,q1−ρq), := K_B,qK_ ,q ND1- _q,~~~~J_q:= K_W,q1- _q,~~~~ _q:= \|W_C^(q) \|_2 (H_q+ L_x,q _x,q1- _q ), (240) Λx,q _x,q :=ρq‖A(q)‖2‖WΔ(q)‖2Hq+‖WB(q)‖2Lx,q(2KΔ,q+Lx,q‖WΔ(q)‖2). := _q \|A^(q) \|_2 \|W_ ^(q) \|_2H_q+ \|W_B^(q) \|_2L_x,q (2K_ ,q+L_x,q \|W_ ^(q) \|_2 ). (241) Proof of Theorem C.3: Fix an arbitrary full parameter tuple W satisfying Assumptions 4.2 and C.1 such that W∈,(Q)W∈S_ γ, δ^(Q). Let ξ∈ξ be some fixed admissible sequence pair of length S (S≥1S≥ 1 is some fixed finite integer). By the definition of a gradient with respect to a tuple in equation (231), it follows that: ‖∇WℒS,(Q)(W)‖F=(∑q=1Q‖∇W(q)ℒS,(Q)(W)‖F2)1/2 \| _WL^S,(Q)(W) \|_F= ( _q=1^Q \| _W^(q)L^S,(Q)(W) \|_F^2 )^1/2 (242) As a result, to derive a bound on the gradient norm with respect to the full model parameter tuple W, it is sufficient to derive the bound on: ‖∇W(q)ℒS,(Q)(W)‖F2 \| _W^(q)L^S,(Q)(W) \|_F^2 (243) for an arbitrary layer q∈1,…,Qq∈\1,…,Q\, and combine the resulting layer-wise bounds. This gradient can be expanded as: ∇W(q)ℒS,(Q)(W)=∇W(q)(1S∑t=1Sℓt(yt(Q)))=1S∑t=1S∇W(q)ℓt(yt(Q)) _W^(q)L^S,(Q)(W)= _W^(q) ( 1S _t=1^S _t (y_t^(Q) ) )= 1S _t=1^S _W^(q) _t (y_t^(Q) ) (244) Applying the chain rule again, we get: vec⁡(∇W(q)ℒS,(Q)(W))=1S∑t=1S(∂yt(Q)∂W(q))⊤∇yℓt(yt(Q)).vec ( _W^(q)L^S,(Q)(W) )= 1S _t=1^S ( ∂ y_t^(Q)∂ W^(q) ) _y _t (y_t^(Q) ). (245) For Q≥2Q≥ 2, we first consider the case q=Q−1q=Q-1 and focus on ∂yt(Q)/∂W(q)∂ y_t^(Q)/∂ W^(q), since ∇yℓt(yt(Q)) _y _t (y_t^(Q) ) is independent of q. By the product rule and using the SSM dynamics, it follows that: ∂yt(Q)∂W(Q−1)=∂(ht(Q))⊤∂W(Q−1)Ct(Q)+(ht(Q))⊤∂Ct(Q)∂W(Q−1) ∂ y_t^(Q)∂ W^(Q-1)= ∂ (h_t^(Q) ) ∂ W^(Q-1)C_t^(Q)+ (h_t^(Q) ) ∂ C_t^(Q)∂ W^(Q-1) (246) Next, we evaluate (246) coordinate-wise, as this avoids any ambiguity associated with calculating the matrix-valued derivatives with respect to the parameter tuples. Let w(Q−1)w^(Q-1) denote an arbitrary scalar coordinate of W(Q−1)W^(Q-1), and let d∈1,…,Dd∈\1,…,D\ denote an arbitrary coordinate of yt(Q)y^(Q)_t. Then: yt,d(Q)=∑n=1Nht,n,d(Q)Ct,n(Q),y_t,d^(Q)= _n=1^Nh_t,n,d^(Q)C_t,n^(Q), (247) and by the product rule: ∂yt,d(Q)∂w(Q−1)=∑n=1N[Ct,n(Q)∂ht,n,d(Q)∂w(Q−1)+ht,n,d(Q)∂Ct,n(Q)∂w(Q−1)] ∂ y_t,d^(Q)∂ w^(Q-1)= _n=1^N [C_t,n^(Q) ∂ h_t,n,d^(Q)∂ w^(Q-1)+h_t,n,d^(Q) ∂ C_t,n^(Q)∂ w^(Q-1) ] (248) Next, since the weight w(Q−1)w^(Q-1) is not a parameter of layer Q, it influences ht(Q)h_t^(Q) at any coordinate only through the inputs to layer Q (i.e. ht(Q)h_t^(Q) is a function of all x1(Q),…,xj(Q),…,xt(Q)x^(Q)_1,…,x^(Q)_j,…,x^(Q)_t, which depend on the parameters of the previous layer through xj(Q)=yj(Q−1)x^(Q)_j=y^(Q-1)_j). As a result: ∂ht,n,d(Q)∂w(Q−1)=∑j=1t∑d′=1D∂ht,n,d(Q)∂xj,d′(Q)∂xj,d′(Q)∂w(Q−1). ∂ h_t,n,d^(Q)∂ w^(Q-1)= _j=1^t _d =1^D ∂ h_t,n,d^(Q)∂ x_j,d ^(Q) ∂ x_j,d ^(Q)∂ w^(Q-1). (249) Furthermore, since Ct,n(Q)=∑d′=1D(WC(Q))n,d′xt,d′(Q)C_t,n^(Q)= _d =1^D (W_C^(Q) )_n,d x_t,d ^(Q), and WC(Q)W_C^(Q) is independent of any parameter of W(Q−1)W^(Q-1), it follows that: ∂Ct,n(Q)∂w(Q−1)=∑d′=1D(WC(Q))n,d′∂xt,d′(Q)∂w(Q−1) ∂ C_t,n^(Q)∂ w^(Q-1)= _d =1^D (W_C^(Q) )_n,d ∂ x_t,d ^(Q)∂ w^(Q-1) (250) Substituting these expressions into (248) results in: ∂yt,d(Q)∂w(Q−1) ∂ y_t,d^(Q)∂ w^(Q-1) =∑n=1N[Ct,n(Q)∑j=1t∑d′=1D∂ht,n,d(Q)∂xj,d′(Q)∂xj,d′(Q)∂w(Q−1)+ht,n,d(Q)∑d′=1D(WC(Q))n,d′∂xt,d′(Q)∂w(Q−1)] = _n=1^N [C_t,n^(Q) _j=1^t _d =1^D ∂ h_t,n,d^(Q)∂ x_j,d ^(Q) ∂ x_j,d ^(Q)∂ w^(Q-1)+h_t,n,d^(Q) _d =1^D (W_C^(Q) )_n,d ∂ x_t,d ^(Q)∂ w^(Q-1) ] =∑n=1N[Ct,n(Q)∑j=1t∑d′=1D∂ht,n,d(Q)∂xj,d′(Q)∂xj,d′(Q)∂w(Q−1)+ht,n,d(Q)∑j=1t∑d′=1Dj=t(WC(Q))n,d′∂xj,d′(Q)∂w(Q−1)] = _n=1^N [C_t,n^(Q) _j=1^t _d =1^D ∂ h_t,n,d^(Q)∂ x_j,d ^(Q) ∂ x_j,d ^(Q)∂ w^(Q-1)+h_t,n,d^(Q) _j=1^t _d =1^D 1_j=t (W_C^(Q) )_n,d ∂ x_j,d ^(Q)∂ w^(Q-1) ] =∑j=1t∑d′=1D[∑n=1N(Ct,n(Q)∂ht,n,d(Q)∂xj,d′(Q)+j=tht,n,d(Q)(WC(Q))n,d′)]∂xj,d′(Q)∂w(Q−1) = _j=1^t _d =1^D [ _n=1^N (C_t,n^(Q) ∂ h_t,n,d^(Q)∂ x_j,d ^(Q)+ 1_j=th_t,n,d^(Q) (W_C^(Q) )_n,d ) ] ∂ x_j,d ^(Q)∂ w^(Q-1) (251) where j=t 1_j=t is 11 if j=tj=t and 00 otherwise. Next, we note that: ∂yt,d(Q)∂xj,d′(Q) ∂ y_t,d^(Q)∂ x_j,d ^(Q) =∂(∑n=1Nht,n,d(Q)Ct,n(Q))∂xj,d′(Q)=∑n=1N∂(ht,n,d(Q)Ct,n(Q))∂xj,d′(Q)=∑n=1N[Ct,n(Q)∂ht,n,d(Q)∂xj,d′(Q)+ht,n,d(Q)∂Ct,n(Q)∂xj,d′(Q)] = ∂ ( _n=1^Nh_t,n,d^(Q)C_t,n^(Q) )∂ x_j,d ^(Q)= _n=1^N ∂ (h_t,n,d^(Q)C_t,n^(Q) )∂ x_j,d ^(Q)= _n=1^N [C_t,n^(Q) ∂ h_t,n,d^(Q)∂ x_j,d ^(Q)+h_t,n,d^(Q) ∂ C_t,n^(Q)∂ x_j,d ^(Q) ] =∑n=1N[Ct,n(Q)∂ht,n,d(Q)∂xj,d′(Q)+ht,n,d(Q)∂(∑d′=1D(WC(Q))n,d′xt,d′(Q))∂xj,d′(Q)] = _n=1^N [C_t,n^(Q) ∂ h_t,n,d^(Q)∂ x_j,d ^(Q)+h_t,n,d^(Q) ∂ ( _d =1^D (W_C^(Q) )_n,d x_t,d ^(Q) )∂ x_j,d ^(Q) ] =∑n=1N(Ct,n(Q)∂ht,n,d(Q)∂xj,d′(Q)+j=tht,n,d(Q)(WC(Q))n,d′) = _n=1^N (C_t,n^(Q) ∂ h_t,n,d^(Q)∂ x_j,d ^(Q)+ 1_j=th_t,n,d^(Q) (W_C^(Q) )_n,d ) (252) where we use the fact that if j≠tj≠ t, ∂(∑d′=1D(WC(Q))n,d′xt,d′(Q))∂xj,d′(Q)=0 ∂ ( _d =1^D (W_C^(Q) )_n,d x_t,d ^(Q) )∂ x_j,d ^(Q)=0 (253) and if j=tj=t, ∂(∑d′=1D(WC(Q))n,d′xj,d′(Q))∂xj,d′(Q)=(WC(Q))n,d′ ∂ ( _d =1^D (W_C^(Q) )_n,d x_j,d ^(Q) )∂ x_j,d ^(Q)= (W_C^(Q) )_n,d (254) As a result, it follows that: ∂yt,d(Q)∂w(Q−1)=∑j=1t∑d′=1D∂yt,d(Q)∂xj,d′(Q)∂xj,d′(Q)∂w(Q−1) ∂ y_t,d^(Q)∂ w^(Q-1)= _j=1^t _d =1^D ∂ y_t,d^(Q)∂ x_j,d ^(Q) ∂ x_j,d ^(Q)∂ w^(Q-1) (255) and finally, since xj,d′(Q)=yj,d′(Q−1)x_j,d ^(Q)=y_j,d ^(Q-1), ∂yt,d(Q)∂w(Q−1)=∑j=1t∑d′=1D∂yt,d(Q)∂xj,d′(Q)∂yj,d′(Q−1)∂w(Q−1) ∂ y_t,d^(Q)∂ w^(Q-1)= _j=1^t _d =1^D ∂ y_t,d^(Q)∂ x_j,d ^(Q) ∂ y_j,d ^(Q-1)∂ w^(Q-1) (256) For Q≥3Q≥ 3, the same coordinate-wise chain rule argument applies to w(Q−2)w^(Q-2). That is, since w(Q−2)w^(Q-2) is not a parameter of layer Q, it influences yt(Q)y^(Q)_t only through the inputs sequence x1(Q),…,xt(Q)x_1^(Q),…,x_t^(Q). As such, it follows that: ∂yt,d(Q)∂w(Q−2)=∑j=1t∑d′=1D∂yt,d(Q)∂xj,d′(Q)∂yj,d′(Q−1)∂w(Q−2) ∂ y_t,d^(Q)∂ w^(Q-2)= _j=1^t _d =1^D ∂ y_t,d^(Q)∂ x_j,d ^(Q) ∂ y_j,d ^(Q-1)∂ w^(Q-2) (257) Furthermore, using the same argument applied to layer Q−1Q-1, we can say that: ∂yj,d′(Q−1)∂w(Q−2)=∑j′=1j∑d′=1D∂yj,d′(Q−1)∂xj′,d′(Q−1)∂yj′,d′(Q−2)∂w(Q−2). ∂ y_j,d ^(Q-1)∂ w^(Q-2)= _j =1^j _d =1^D ∂ y_j,d ^(Q-1)∂ x_j ,d ^(Q-1) ∂ y_j ,d ^(Q-2)∂ w^(Q-2). (258) Substituting this result into (257), we obtain: ∂yt,d(Q)∂w(Q−2)=∑j=1t∑d′=1D∂yt,d(Q)∂xj,d′(Q)∑j′=1j∑d′=1D∂yj,d′(Q−1)∂xj′,d′(Q−1)∂yj′,d′(Q−2)∂w(Q−2)=∑j=1t∑j′=1j∑d′=1D∑d′=1D∂yt,d(Q)∂xj,d′(Q)∂yj,d′(Q−1)∂xj′,d′(Q−1)∂yj′,d′(Q−2)∂w(Q−2) ∂ y_t,d^(Q)∂ w^(Q-2)= _j=1^t _d =1^D ∂ y_t,d^(Q)∂ x_j,d ^(Q) _j =1^j _d =1^D ∂ y_j,d ^(Q-1)∂ x_j ,d ^(Q-1) ∂ y_j ,d ^(Q-2)∂ w^(Q-2)= _j=1^t _j =1^j _d =1^D _d =1^D ∂ y_t,d^(Q)∂ x_j,d ^(Q) ∂ y_j,d ^(Q-1)∂ x_j ,d ^(Q-1) ∂ y_j ,d ^(Q-2)∂ w^(Q-2) (259) This result means that, compared to the case of w(Q−1)w^(Q-1), the derivative with respect to w(Q−2)w^(Q-2) contains one additional input-output derivative, which corresponds to its propagation through layer Q−1Q-1. Applying this argument recursively, we derive the final derivative of the d-th coordinate of the final output with respect to some parameter w(q)w^(q) of the layer q∈1,…,Q−1q∈\1,…,Q-1\ as: ∂yt,d(Q)∂w(q)=∑jQ−1=1jQ⋯∑jq=1jq+1∑dQ−1=1D⋯∑dq=1D[∏q′=q+1Q∂yjq′,dq′(q′)∂xjq′−1,dq′−1(q′)]∂yjq,dq(q)∂w(q), ∂ y_t,d^(Q)∂ w^(q)= _j_Q-1=1^j_Q·s _j_q=1^j_q+1 _d_Q-1=1^D·s _d_q=1^D [ _q =q+1^Q ∂ y_j_q ,d_q ^(q )∂ x_j_q -1,d_q -1^(q ) ] ∂ y_j_q,d_q^(q)∂ w^(q), (260) where we also have jQ=tj_Q=t, dQ=d_Q=d, and jq≤jq+1≤⋯≤jQ−1≤jQ=tj_q≤ j_q+1≤·s≤ j_Q-1≤ j_Q=t. We next collect the coordinate derivatives in (260) into matrices, which transforms sums over feature coordinates into matrix products, while the sums over time indices remain explicit. For each layer q, define the input-output Jacobian matrix as: t,j(q):=∂yt(q)∂xj(q)=(∂yt,1(q)∂xj,1(q)⋯∂yt,1(q)∂xj,D(q)⋱∂yt,D(q)∂xj,1(q)…∂yt,D(q)∂xj,D(q))∈ℝD×D.K_t,j^(q):= ∂ y_t^(q)∂ x_j^(q)= pmatrix ∂ y_t,1^(q)∂ x_j,1^(q)&·s& ∂ y_t,1^(q)∂ x_j,D^(q)\\ & & \\ ∂ y_t,D^(q)∂ x_j,1^(q)&…& ∂ y_t,D^(q)∂ x_j,D^(q) pmatrix ^D× D. (261) The Jacobian is computed with the parameters of layer q and all input tokens other than xj(q)x_j^(q) held fixed, and evaluated at the input sequence received by that layer. Additionally, by causality, t,j(q)=D×DK_t,j^(q)=0_D× D for j>tj>t. Then, for any layer q∈1,…,Q−1q∈\1,…,Q-1\, for any subsequent layer q′∈q+1,…,Qq ∈\q+1,…,Q\, the parameters W(q)W^(q) affect y(q′)y^(q ) only through the input sequence to layer q′q . Since xj(q′)=yj(q′−1)x_j^(q )=y_j^(q -1), the chain rule can be written in the matrix form as: ∂yt(q′)∂W(q)=∑j=1t∂yt(q′)∂xj(q′)∂xj(q′)∂W(q)=∑j=1t,j(q′)∂yj(q′−1)∂W(q). ∂ y_t^(q )∂ W^(q)= _j=1^t ∂ y_t^(q )∂ x_j^(q ) ∂ x_j^(q )∂ W^(q)= _j=1^tK_t,j^(q ) ∂ y_j^(q -1)∂ W^(q). (262) Applying this recursively, we get: ∂yt(Q)∂W(q)=∑jQ−1=1jQ⋯∑jq=1jq+1[jQ,jQ−1(Q)×⋯×jq+1,jq(q+1)]∂yjq(q)∂W(q) ∂ y_t^(Q)∂ W^(q)= _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [K_j_Q,j_Q-1^(Q)×·s×K_j_q+1,j_q^(q+1) ] ∂ y_j_q^(q)∂ W^(q) (263) for any q∈1,…,Q−1q∈\1,…,Q-1\, the matrix product contains only t,jQ−1(Q)K_t,j_Q-1^(Q) for the case q=Q−1q=Q-1, the × operator denotes matrix multiplication, and jQ=tj_Q=t. Note that equation (263) provides an explicit expression for the parameter Jacobian ∂yt(Q)∂W(q) ∂ y_t^(Q)∂ W^(q) appearing in (245), in terms of the own-layer parameter Jacobians ∂yjq(q)∂W(q) ∂ y_j_q^(q)∂ W^(q) and the input-output Jacobians of the subsequent layers. We next use this decomposition to bound the norm of the parameter Jacobian. We use the Frobenius norm ∥⋅∥F\|·\|_F for the flattened parameter Jacobians and the spectral norm ∥⋅∥2\|·\|_2 (where for M∈ℝD×DM ^D× D, ‖M‖2:=supv∈ℝD‖v‖2=1‖Mv‖2\|M\|_2:= _ subarraycv ^D\\ \|v\|_2=1 subarray\|Mv\|_2) for the input-output matrices. Applying the triangle inequality and the fact that ‖AB‖F≤‖A‖2‖B‖F\|AB\|_F≤\|A\|_2\|B\|_F, it follows that: ‖∂yt(Q)∂W(q)‖F \| ∂ y_t^(Q)∂ W^(q) \|_F =‖∑jQ−1=1jQ⋯∑jq=1jq+1[jQ,jQ−1(Q)×⋯×jq+1,jq(q+1)]∂yjq(q)∂W(q)‖F = \| _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [K_j_Q,j_Q-1^(Q)×·s×K_j_q+1,j_q^(q+1) ] ∂ y_j_q^(q)∂ W^(q) \|_F ≤∑jQ−1=1jQ⋯∑jq=1jq+1‖[jQ,jQ−1(Q)×⋯×jq+1,jq(q+1)]∂yjq(q)∂W(q)‖F ≤ _j_Q-1=1^j_Q·s _j_q=1^j_q+1 \| [K_j_Q,j_Q-1^(Q)×·s×K_j_q+1,j_q^(q+1) ] ∂ y_j_q^(q)∂ W^(q) \|_F ≤∑jQ−1=1jQ⋯∑jq=1jq+1[∏q′=q+1Q‖jq′,jq′−1(q′)‖2]‖∂yjq(q)∂W(q)‖F, ≤ _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [ _q =q+1^Q \|K_j_q ,j_q -1^(q ) \|_2 ] \| ∂ y_j_q^(q)∂ W^(q) \|_F, (264) where jQ=tj_Q=t and q∈1,…,Q−1q∈\1,…,Q-1\. Next, we bound the terms of this expression separately. Thanks to Assumption C.1 and Definition C.2, we can apply the single-layer gradient proof argument from Appendix A to each layer q, bounding the term ‖∂yjq(q)∂W(q)‖F \| ∂ y_j_q^(q)∂ W^(q) \|_F. First, for each q∈1,…,Qq∈\1,…,Q\, let: KΔ,q K_ ,q :=‖WΔ(q)‖2Lx,q+Dlog⁡(2),KB,q:=‖WB(q)‖2Lx,q2, := \|W_ ^(q) \|_2L_x,q+ D (2),~~~~K_B,q:= \|W_B^(q) \|_2L_x,q^2, (265) KW,q K_W,q :=NDKB,qKΔ,q(KΔ,q+D‖A(q)‖2Lx,q)(ρq1−ρq)+Lx,q2NKΔ,q+Lx,qDKB,q. :=NDK_B,qK_ ,q (K_ ,q+ D \|A^(q) \|_2L_x,q ) ( _q1- _q )+L_x,q^2 NK_ ,q+L_x,q DK_B,q. (266) Given these definitions, we apply the single-layer argument from Appendix A for every layer q∈1,…,Qq∈\1,…,Q\. Since the input to any layer q depends only on the parameters of the preceding layers, it remains fixed when differentiating with respect to W(q)W^(q). It therefore follows that: ‖hj(q)‖F \|h_j^(q) \|_F ≤KB,qKΔ,qND1−ρq:=Hq, ≤ K_B,qK_ ,q ND1- _q:=H_q, (267) ‖∂hj(q)∂W(q)‖F \| ∂ h_j^(q)∂ W^(q) \|_F ≤KW,q1−ρq:=Jq, ≤ K_W,q1- _q:=J_q, (268) for every q∈1,…,Qq∈\1,…,Q\ and j∈1,…,Sj∈\1,…,S\. Here, the state Jacobian is flattened as in the single-layer analysis. Then, by definition of Cj(q)C_j^(q), Assumption C.1 gives: ‖Cj(q)‖2=‖WC(q)xj(q)‖2≤‖WC(q)‖2Lx,q. \|C_j^(q) \|_2= \|W_C^(q)x_j^(q) \|_2≤ \|W_C^(q) \|_2L_x,q. (269) Finally, by applying the same product-rule argument to yj(q)=(hj(q))⊤Cj(q)y_j^(q)= (h_j^(q) ) C_j^(q) as in Appendix A, and using (267) and (268), we arrive at: ‖∂yj(q)∂W(q)‖F \| ∂ y_j^(q)∂ W^(q) \|_F ≤‖Cj(q)‖2‖∂hj(q)∂W(q)‖F+‖hj(q)‖F‖xj(q)‖2 ≤ \|C_j^(q) \|_2 \| ∂ h_j^(q)∂ W^(q) \|_F+ \|h_j^(q) \|_F \|x_j^(q) \|_2 ≤‖WC(q)‖2Lx,qJq+HqLx,q=Lx,q(Hq+‖WC(q)‖2Jq). ≤ \|W_C^(q) \|_2L_x,qJ_q+H_qL_x,q=L_x,q (H_q+ \|W_C^(q) \|_2J_q ). (270) This bound holds for every q∈1,…,Qq∈\1,…,Q\ and j∈1,…,Sj∈\1,…,S\ uniformly over all admissible input sequences. Substituting (270) into (264), we get: ‖∂yt(Q)∂W(q)‖F≤Lx,q(Hq+‖WC(q)‖2Jq)(∑jQ−1=1jQ⋯∑jq=1jq+1[∏q′=q+1Q‖jq′,jq′−1(q′)‖2]) \| ∂ y_t^(Q)∂ W^(q) \|_F≤ L_x,q (H_q+ \|W_C^(q) \|_2J_q ) ( _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [ _q =q+1^Q \|K_j_q ,j_q -1^(q ) \|_2 ] ) (271) Next, we derive a bound on the term: ∑jQ−1=1jQ⋯∑jq=1jq+1[∏q′=q+1Q‖jq′,jq′−1(q′)‖2] _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [ _q =q+1^Q \|K_j_q ,j_q -1^(q ) \|_2 ] (272) To derive the bound, for an arbitrary layer q′∈q+1,…,Qq ∈\q+1,…,Q\, we begin by expressing t,j(q′)K_t,j^(q ) in terms of the layer’s hidden-state sensitivity. Then, by definition: (t,j(q′))d,d′=∂yt,d(q′)∂xj,d′(q′)=∑n=1NCt,n(q′)∂ht,n,d(q′)∂xj,d′(q′)+j=t∑n=1Nht,n,d(q′)(WC(q′))n,d′, (K_t,j^(q ) )_d,d = ∂ y_t,d^(q )∂ x_j,d ^(q )= _n=1^NC_t,n^(q ) ∂ h_t,n,d^(q )∂ x_j,d ^(q )+ 1_j=t _n=1^Nh_t,n,d^(q ) (W_C^(q ) )_n,d , (273) for 1≤j≤t≤S1≤ j≤ t≤ S and d,d′∈1,…,Dd,d ∈\1,…,D\. Next, by definition of the spectral norm, we know that: ‖t,j(q′)‖2=supu∈ℝD‖u‖2=1‖t,j(q′)u‖2. \|K_t,j^(q ) \|_2= _ subarraycu ^D\\ \|u\|_2=1 subarray \|K_t,j^(q )u \|_2. (274) As such, we first evaluate the value of t,j(q′)uK_t,j^(q )u for some arbitrary unit vector u∈ℝDu ^D. Multiplying (273) by ud′u_d , summing over d′d , and collecting the output coordinates, we get: t,j(q′)u=(∑d′=1Dud′∂ht(q′)∂xj,d′(q′))⊤Ct(q′)+j=t(ht(q′))⊤WC(q′)u.K_t,j^(q )u= ( _d =1^Du_d ∂ h_t^(q )∂ x_j,d ^(q ) ) C_t^(q )+ 1_j=t (h_t^(q ) ) W_C^(q )u. (275) Note that, since u is a unit vector, the sum in the first term is, by definition, the directional derivative of the hidden state with respect to xj(q′)x_j^(q ) in direction u. We denote it as: Dxj(q′)ht(q′)[u]:=∑d′=1Dud′∂ht(q′)∂xj,d′(q′)∈ℝN×D.D_x_j^(q )h_t^(q )[u]:= _d =1^Du_d ∂ h_t^(q )∂ x_j,d ^(q ) ^N× D. (276) As such, it follows that: t,j(q′)u=(Dxj(q′)ht(q′)[u])⊤Ct(q′)+j=t(ht(q′))⊤WC(q′)uK_t,j^(q )u= (D_x_j^(q )h_t^(q )[u] ) C_t^(q )+ 1_j=t (h_t^(q ) ) W_C^(q )u (277) We then bound both terms of the sum in (277) separately. We evaluate the term (Dxj(q′)ht(q′)[u])⊤Ct(q′) (D_x_j^(q )h_t^(q )[u] ) C_t^(q ) first. Recall that by (227): ht(q′)=A¯t(q′)⊙ht−1(q′)+B¯t(q′)diag⁡(xt(q′)).h_t^(q )= A_t^(q ) h_t-1^(q )+ B_t^(q )diag (x_t^(q ) ). Since, for fixed layer q′q weights, A¯t(q′) A_t^(q ) and B¯t(q′) B_t^(q ) depend only on the input xt(q′)x_t^(q ), their derivatives with respect to xj(q′)x_j^(q ) vanish when t≠jt≠ j. Furthermore, Dxj(q′)xt(q′)[u]=j=tuD_x_j^(q )x_t^(q )[u]= 1_j=tu. As such, applying these results and the product rule, we get: Dxj(q′)ht(q′)[u]=A¯t(q′)⊙Dxj(q′)ht−1(q′)[u] D_x_j^(q )h_t^(q )[u]= A_t^(q ) D_x_j^(q )h_t-1^(q )[u] +j=t[(Dxt(q′)A¯t(q′)[u])⊙ht−1(q′)+ + 1_j=t [ (D_x_t^(q ) A_t^(q )[u] ) h_t-1^(q )+ (Dxt(q′)B¯t(q′)[u])diag(xt(q′))+B¯t(q′)diag(u)], (D_x_t^(q ) B_t^(q )[u] )diag (x_t^(q ) )+ B_t^(q )diag(u) ], (278) where we also used the fact that Dxj(q′)diag⁡(xt(q′))[u]=j=tdiag⁡(u)D_x_j^(q )diag (x_t^(q ) )[u]= 1_j=tdiag (u ). Next, we consider three cases. When 0≤t<j0≤ t<j, because the hidden state is initialized to 00 and its evolution is causal, then: Dxj(q′)ht(q′)[u]=,∀t∈0,…,j−1.D_x_j^(q )h_t^(q )[u]=0, ∀ t∈\0,…,j-1\. (279) When, t=jt=j, Dxj(q′)ht−1(q′)[u]=Dxt(q′)ht−1(q′)[u]=D_x_j^(q )h_t-1^(q )[u]=D_x_t^(q )h_t-1^(q )[u]=0, so: Dxt(q′)ht(q′)[u]=(Dxt(q′)A¯t(q′)[u])⊙ht−1(q′)+(Dxt(q′)B¯t(q′)[u])diag⁡(xt(q′))+B¯t(q′)diag⁡(u).D_x_t^(q )h_t^(q )[u]= (D_x_t^(q ) A_t^(q )[u] ) h_t-1^(q )+ (D_x_t^(q ) B_t^(q )[u] )diag (x_t^(q ) )+ B_t^(q )diag(u). (280) Finally, for t>jt>j, the recurrence becomes: Dxj(q′)ht(q′)[u]=A¯t(q′)⊙Dxj(q′)ht−1(q′)[u].D_x_j^(q )h_t^(q )[u]= A_t^(q ) D_x_j^(q )h_t-1^(q )[u]. (281) Next, by equation (236), which follows from definition C.2, we obtain: ‖Dxj(q′)ht(q′)[u]‖F≤ρq′‖Dxj(q′)ht−1(q′)[u]‖F≤ρq′t−j‖Dxj(q′)hj(q′)[u]‖F,∀t>j. \|D_x_j^(q )h_t^(q )[u] \|_F≤ _q \|D_x_j^(q )h_t-1^(q )[u] \|_F≤ _q ^t-j \|D_x_j^(q )h_j^(q )[u] \|_F, ∀ t>j. (282) As a result, we just need to bound ‖Dxj(q′)ht(q′)[u]‖F \|D_x_j^(q )h_t^(q )[u] \|_F for the case when t=jt=j, and apply the recursive relationship. By the triangle inequality, when t=jt=j, it follows that: ‖Dxj(q′)hj(q′)[u]‖F≤ \|D_x_j^(q )h_j^(q )[u] \|_F≤ ‖(Dxj(q′)A¯j(q′)[u])⊙hj−1(q′)‖F \| (D_x_j^(q ) A_j^(q )[u] ) h_j-1^(q ) \|_F +‖(Dxj(q′)B¯j(q′)[u])diag⁡(xj(q′))‖F+‖B¯j(q′)diag⁡(u)‖F + \| (D_x_j^(q ) B_j^(q )[u] )diag (x_j^(q ) ) \|_F+ \| B_j^(q )diag(u) \|_F (283) Then, evaluating it term by term Dxj(q′)A¯j(q′)[u] D_x_j^(q ) A_j^(q )[u] =Dxj(q′)(exp⁡(A(q′)(Δj(q′))⊤))[u] =D_x_j^(q ) ( (A^(q ) ( _j^(q ) ) ) )[u] =exp⁡(A(q′)(Δj(q′))⊤)⊙[A(q′)(Dxj(q′)Δj(q′)[u])⊤] = (A^(q ) ( _j^(q ) ) ) [A^(q ) (D_x_j^(q ) _j^(q )[u] ) ] =A¯j(q′)⊙[A(q′)(Dxj(q′)Δj(q′)[u])⊤]. = A_j^(q ) [A^(q ) (D_x_j^(q ) _j^(q )[u] ) ]. (284) Using the definition: Dxj(q′)Δj(q′)[u]=Dxj(q′)(Softplus⁡(WΔ(q′)xj(q′)))[u]=Softplus′⁡(WΔ(q′)xj(q′))⊙(WΔ(q′)u). D_x_j^(q ) _j^(q )[u]=D_x_j^(q ) (Softplus (W_ ^(q )x_j^(q ) ) )[u]=Softplus (W_ ^(q )x_j^(q ) ) (W_ ^(q )u ). (285) Next, since Softplus′⁡(z)=11+e−z∈(0,1)Softplus (z)= 11+e^-z∈(0,1), we get: ‖Dxj(q′)Δj(q′)[u]‖2≤‖WΔ(q′)‖2‖u‖2=‖WΔ(q′)‖2. \|D_x_j^(q ) _j^(q )[u] \|_2≤ \|W_ ^(q ) \|_2\|u\|_2= \|W_ ^(q ) \|_2. (286) Then, the fact that ‖ab⊤‖F=‖a‖2‖b‖2\|ab \|_F=\|a\|_2\|b\|_2 (which follows from the definition of the Frobenius norm), the following holds: ‖Dxj(q′)A¯j(q′)[u]‖F \|D_x_j^(q ) A_j^(q )[u] \|_F ≤‖A¯j(q′)‖max‖A(q′)(Dxj(q′)Δj(q′)[u])⊤‖F ≤ \| A_j^(q ) \|_ \|A^(q ) (D_x_j^(q ) _j^(q )[u] ) \|_F =‖A¯j(q′)‖max‖A(q′)‖2‖Dxj(q′)Δj(q′)[u]‖2≤ρq′‖A(q′)‖2‖WΔ(q′)‖2, = \| A_j^(q ) \|_ \|A^(q ) \|_2 \|D_x_j^(q ) _j^(q )[u] \|_2≤ _q \|A^(q ) \|_2 \|W_ ^(q ) \|_2, (287) Finally, combining equations (287) and (267) yields: ‖(Dxj(q′)A¯j(q′)[u])⊙hj−1(q′)‖F≤‖Dxj(q′)A¯j(q′)[u]‖F‖hj−1(q′)‖F≤ρq′‖A(q′)‖2‖WΔ(q′)‖2Hq′. \| (D_x_j^(q ) A_j^(q )[u] ) h_j-1^(q ) \|_F≤ \|D_x_j^(q ) A_j^(q )[u] \|_F \|h_j-1^(q ) \|_F≤ _q \|A^(q ) \|_2 \|W_ ^(q ) \|_2H_q . (288) For the second term, from the definition of B¯j(q′) B_j^(q ), and using the product rule, we have: Dxj(q′)B¯j(q′)[u]=(WB(q′)u)(Δj(q′))⊤+(WB(q′)xj(q′))(Dxj(q′)Δj(q′)[u])⊤D_x_j^(q ) B_j^(q )[u]= (W_B^(q )u ) ( _j^(q ) ) + (W_B^(q )x_j^(q ) ) (D_x_j^(q ) _j^(q )[u] ) (289) Then, combining the result of (286) with ‖Δj(q′)‖2≤‖WΔ(q′)xj(q′)‖2+Dlog⁡(2)≤KΔ,q′, \| _j^(q ) \|_2≤ \|W_ ^(q )x_j^(q ) \|_2+ D (2)≤ K_ ,q , (290) we get: ‖Dxj(q′)B¯j(q′)[u]‖F \|D_x_j^(q ) B_j^(q )[u] \|_F ≤‖WB(q′)u‖2‖Δj(q′)‖2+‖WB(q′)xj(q′)‖2‖Dxj(q′)Δj(q′)[u]‖2 ≤ \|W_B^(q )u \|_2 \| _j^(q ) \|_2+ \|W_B^(q )x_j^(q ) \|_2 \|D_x_j^(q ) _j^(q )[u] \|_2 ≤‖WB(q′)‖2KΔ,q′+‖WB(q′)‖2Lx,q′‖WΔ(q′)‖2. ≤ \|W_B^(q ) \|_2K_ ,q + \|W_B^(q ) \|_2L_x,q \|W_ ^(q ) \|_2. (291) Finally, since ‖diag⁡(v)‖2=maxd⁡|vd|=‖v‖∞≤‖v‖2 \|diag (v ) \|_2= _d |v_d |= \|v \|_∞≤ \|v \|_2 (and therefore ‖diag⁡(xj(q′))‖2≤Lx,q′ \|diag (x_j^(q ) ) \|_2≤ L_x,q ), the second term is bounded as: ‖(Dxj(q′)B¯j(q′)[u])diag⁡(xj(q′))‖F \| (D_x_j^(q ) B_j^(q )[u] )diag (x_j^(q ) ) \|_F ≤‖Dxj(q′)B¯j(q′)[u]‖F‖diag⁡(xj(q′))‖2 ≤ \|D_x_j^(q ) B_j^(q )[u] \|_F \|diag (x_j^(q ) ) \|_2 ≤‖WB(q′)‖2Lx,q′(KΔ,q′+Lx,q′‖WΔ(q′)‖2) ≤ \|W_B^(q ) \|_2L_x,q (K_ ,q +L_x,q \|W_ ^(q ) \|_2 ) (292) For the remaining term, the bound is as follows: ‖B¯j(q′)diag⁡(u)‖F \| B_j^(q )diag(u) \|_F ≤‖B¯j(q′)‖F‖diag⁡(u)‖2 ≤ \| B_j^(q ) \|_F \|diag(u) \|_2 ≤‖WB(q′)xj(q′)‖2‖Δj(q′)‖2‖u‖2 ≤ \|W_B^(q )x_j^(q ) \|_2 \| _j^(q ) \|_2 \|u \|_2 ≤‖WB(q′)‖2Lx,q′KΔ,q′ ≤ \|W_B^(q ) \|_2L_x,q K_ ,q (293) Combining the three bounds, we get: ‖Dxj(q′)hj(q′)[u]‖F≤ρq′‖A(q′)‖2‖WΔ(q′)‖2Hq′+‖WB(q′)‖2Lx,q′(2KΔ,q′+Lx,q′‖WΔ(q′)‖2). \|D_x_j^(q )h_j^(q )[u] \|_F≤ _q \|A^(q ) \|_2 \|W_ ^(q ) \|_2H_q + \|W_B^(q ) \|_2L_x,q (2K_ ,q +L_x,q \|W_ ^(q ) \|_2 ). (294) For ease of notation, for each layer q′q , we define: Λx,q′:=ρq′‖A(q′)‖2‖WΔ(q′)‖2Hq′+‖WB(q′)‖2Lx,q′(2KΔ,q′+Lx,q′‖WΔ(q′)‖2). _x,q := _q \|A^(q ) \|_2 \|W_ ^(q ) \|_2H_q + \|W_B^(q ) \|_2L_x,q (2K_ ,q +L_x,q \|W_ ^(q ) \|_2 ). (295) Using this notation, we have the following bound: ‖Dxj(q′)hj(q′)[u]‖F≤Λx,q′,‖u‖2=1. \|D_x_j^(q )h_j^(q )[u] \|_F≤ _x,q , \|u\|_2=1. (296) Substituting this bound into the recurrence established in (282), we get: ‖Dxj(q′)ht(q′)[u]‖F≤ρq′t−j‖Dxj(q′)hj(q′)[u]‖F≤Λx,q′ρq′t−j,1≤j≤t≤S. \|D_x_j^(q )h_t^(q )[u] \|_F≤ _q ^t-j \|D_x_j^(q )h_j^(q )[u] \|_F≤ _x,q _q ^t-j, 1≤ j≤ t≤ S. (297) given ‖u‖2=1\|u\|_2=1, and ‖Dxj(q′)ht(q′)[u]‖F=0 \|D_x_j^(q )h_t^(q )[u] \|_F=0 for t<jt<j. Returning to the expression for t,j(q′)uK_t,j^(q )u in (277), we can bound the first term of the sum using the bound derived in (297) (and the fact that ‖Z⊤c‖2≤‖Z‖F‖c‖2 \|Z c \|_2≤\|Z\|_F\|c\|_2) as follows: ‖(Dxj(q′)ht(q′)[u])⊤Ct(q′)‖2≤‖Dxj(q′)ht(q′)[u]‖F‖Ct(q′)‖2≤‖WC(q′)‖2Lx,q′Λx,q′ρq′t−j. \| (D_x_j^(q )h_t^(q )[u] ) C_t^(q ) \|_2≤ \|D_x_j^(q )h_t^(q )[u] \|_F \|C_t^(q ) \|_2≤ \|W_C^(q ) \|_2L_x,q _x,q _q ^t-j. (298) For the second term, we can bound this as: ‖j=t(ht(q′))⊤WC(q′)u‖2≤j=t‖ht(q′)‖F‖WC(q′)‖2‖u‖2=j=t‖ht(q′)‖F‖WC(q′)‖2 \| 1_j=t (h_t^(q ) ) W_C^(q )u \|_2≤ 1_j=t \|h_t^(q ) \|_F \|W_C^(q ) \|_2 \|u \|_2= 1_j=t \|h_t^(q ) \|_F \|W_C^(q ) \|_2 (299) Using the triangle inequality, and combining the derived bounds on the terms of the sum in (277), we get: ‖t,j(q′)u‖2≤‖WC(q′)‖2Lx,q′Λx,q′ρq′t−j+j=t‖WC(q′)‖2Hq′. \|K_t,j^(q )u \|_2≤ \|W_C^(q ) \|_2L_x,q _x,q _q ^t-j+ 1_j=t \|W_C^(q ) \|_2H_q . (300) Since the right hand side of (300) is independent of u, and therefore holds uniformly for all unit vectors u, taking the supremum over all u, we get: ‖t,j(q′)‖2≤‖WC(q′)‖2Lx,q′Λx,q′ρq′t−j+j=t‖WC(q′)‖2Hq′,1≤j≤t≤S. \|K_t,j^(q ) \|_2≤ \|W_C^(q ) \|_2L_x,q _x,q _q ^t-j+ 1_j=t \|W_C^(q ) \|_2H_q , 1≤ j≤ t≤ S. (301) Next, for every fixed position t, summing ‖t,j(q′)‖2 \|K_t,j^(q ) \|_2 over j yields: ∑j=1t‖t,j(q′)‖2 _j=1^t \|K_t,j^(q ) \|_2 ≤∑j=1t(‖WC(q′)‖2Lx,q′Λx,q′ρq′t−j+j=t‖WC(q′)‖2Hq′) ≤ _j=1^t ( \|W_C^(q ) \|_2L_x,q _x,q _q ^t-j+ 1_j=t \|W_C^(q ) \|_2H_q ) =‖WC(q′)‖2Lx,q′Λx,q′∑j=1tρq′t−j+‖WC(q′)‖2Hq′ = \|W_C^(q ) \|_2L_x,q _x,q _j=1^t _q ^t-j+ \|W_C^(q ) \|_2H_q =‖WC(q′)‖2Lx,q′Λx,q′1−ρq′t1−ρq′+‖WC(q′)‖2Hq′ = \|W_C^(q ) \|_2L_x,q _x,q 1- _q ^t1- _q + \|W_C^(q ) \|_2H_q ≤‖WC(q′)‖2(Hq′+Lx,q′Λx,q′1−ρq′)=κq′, ≤ \|W_C^(q ) \|_2 (H_q + L_x,q _x,q 1- _q )= _q , (302) where, for each layer q′q , we define the uniform layerwise input-output sensitivity bound κq′:=‖WC(q′)‖2(Hq′+Lx,q′Λx,q′1−ρq′) _q := \|W_C^(q ) \|_2 (H_q + L_x,q _x,q 1- _q ) (303) Next, since κq′ _q is independent of t, it follows that: max⁡∑j=1t1≤t≤S⁡‖t,j(q′)‖2≤κq′. _1≤ t≤ S _j=1^t \|K_t,j^(q ) \|_2≤ _q . (304) Since W satisfies the stated stability and bounded-input conditions, the bound holds uniformly over all admissible input-target sequences and token indices. Now, returning to (272), consider the two inner-most sums (for q≤Q−2q≤ Q-2): ∑jq+1=1jq+2∑jq=1jq+1[∏q′=q+1Q‖jq′,jq′−1(q′)‖2] _j_q+1=1^j_q+2 _j_q=1^j_q+1 [ _q =q+1^Q \|K_j_q ,j_q -1^(q ) \|_2 ] =∑jq+1=1jq+2∑jq=1jq+1‖jq+1,jq(q+1)‖2‖jq+2,jq+1(q+2)‖2⋯‖jQ,jQ−1(Q)‖2 = _j_q+1=1^j_q+2 _j_q=1^j_q+1 \|K_j_q+1,j_q^(q+1) \|_2 \|K_j_q+2,j_q+1^(q+2) \|_2·s \|K_j_Q,j_Q-1^(Q) \|_2 =∑jq+1=1jq+2[∑jq=1jq+1‖jq+1,jq(q+1)‖2‖jq+2,jq+1(q+2)‖2⋯‖jQ,jQ−1(Q)‖2] = _j_q+1=1^j_q+2 [ _j_q=1^j_q+1 \|K_j_q+1,j_q^(q+1) \|_2 \|K_j_q+2,j_q+1^(q+2) \|_2·s \|K_j_Q,j_Q-1^(Q) \|_2 ] =∑jq+1=1jq+2‖jq+2,jq+1(q+2)‖2⋯‖jQ,jQ−1(Q)‖2[∑jq=1jq+1‖jq+1,jq(q+1)‖2] = _j_q+1=1^j_q+2 \|K_j_q+2,j_q+1^(q+2) \|_2·s \|K_j_Q,j_Q-1^(Q) \|_2 [ _j_q=1^j_q+1 \|K_j_q+1,j_q^(q+1) \|_2 ] ≤∑jq+1=1jq+2‖jq+2,jq+1(q+2)‖2⋯‖jQ,jQ−1(Q)‖2κq+1 ≤ _j_q+1=1^j_q+2 \|K_j_q+2,j_q+1^(q+2) \|_2·s \|K_j_Q,j_Q-1^(Q) \|_2 _q+1 =κq+1∑jq+1=1jq+2‖jq+2,jq+1(q+2)‖2⋯‖jQ,jQ−1(Q)‖2, = _q+1 _j_q+1=1^j_q+2 \|K_j_q+2,j_q+1^(q+2) \|_2·s \|K_j_Q,j_Q-1^(Q) \|_2, (305) where the inequality follows from (304). Since κq+1 _q+1 is independent of all remaining sums, and we can apply the same argument successively to each remaining innermost sum, it follows that: ∑jQ−1=1jQ⋯∑jq=1jq+1[∏q′=q+1Q‖jq′,jq′−1(q′)‖2]≤κq+1κq+2⋯κQ=∏q′=q+1Qκq′. _j_Q-1=1^j_Q·s _j_q=1^j_q+1 [ _q =q+1^Q \|K_j_q ,j_q -1^(q ) \|_2 ]≤ _q+1 _q+2·s _Q= _q =q+1^Q _q . (306) for the fixed W and every admissible input-target sequence pair. For q=Q−1q=Q-1, there is only one sum, and holds directly from definition: ∑jQ−1=1t‖t,jQ−1(Q)‖2≤κQ _j_Q-1=1^t \|K_t,j_Q-1^(Q) \|_2≤ _Q. Finally, combining the nested-sum bound in (306) with (264) and (270), we obtain: ‖∂yt(Q)∂W(q)‖F≤Lx,q(Hq+‖WC(q)‖2Jq)∏q′=q+1Qκq′. \| ∂ y_t^(Q)∂ W^(q) \|_F≤ L_x,q (H_q+ \|W_C^(q) \|_2J_q ) _q =q+1^Q _q . (307) This bound holds for every q∈1,…,Qq∈\1,…,Q\ and t∈1,…,St∈\1,…,S\, where an empty product is interpreted as one. For the q=Qq=Q case, the inequality follows directly from the own-layer bound in (270). We now return to the loss gradient in (245). Under Assumption 4.2, applied at the fixed full parameter tuple W to final-layer predictions uniformly over admissible pairs and token indices, we have: ‖∇yℓt(yt(Q))‖2≤Lℓ. \| _y _t (y_t^(Q) ) \|_2≤ L_ . (308) Finally, taking the Euclidean norm of the vectorized gradient in (245), which equals its block Frobenius norm, gives: ‖∇W(q)ℒS,(Q)(W)‖F \| _W^(q)L^S,(Q)(W) \|_F =‖1S∑t=1S(∂yt(Q)∂W(q))⊤∇yℓt(yt(Q))‖2 = \| 1S _t=1^S ( ∂ y_t^(Q)∂ W^(q) ) _y _t (y_t^(Q) ) \|_2 ≤1S∑t=1S‖∇yℓt(yt(Q))‖2‖∂yt(Q)∂W(q)‖F ≤ 1S _t=1^S \| _y _t (y_t^(Q) ) \|_2 \| ∂ y_t^(Q)∂ W^(q) \|_F ≤1S∑t=1SLℓLx,q(Hq+‖WC(q)‖2Jq)∏q′=q+1Qκq′ ≤ 1S _t=1^SL_ L_x,q (H_q+ \|W_C^(q) \|_2J_q ) _q =q+1^Q _q =LℓLx,q(Hq+‖WC(q)‖2Jq)∏q′=q+1Qκq′, =L_ L_x,q (H_q+ \|W_C^(q) \|_2J_q ) _q =q+1^Q _q , (309) where we use (308), (307), the triangle inequality, and the fact that ‖M⊤v‖2≤‖M‖F‖v‖2\|M v\|_2≤\|M\|_F\|v\|_2. The bound is independent of the token index t. Substituting the result of (309) back into (242), we get: ‖∇WℒS,(Q)(W)‖F≤Lℓ(∑q=1QLx,q2(Hq+‖WC(q)‖2Jq)2∏q′=q+1Qκq′2)1/2=SSM(Q)(W). \| _WL^S,(Q)(W) \|_F≤ L_ ( _q=1^QL_x,q^2 (H_q+ \|W_C^(q) \|_2J_q )^2 _q =q+1^Q _q ^2 )^1/2=G_SSM^(Q)(W). (310) Since S≥1S≥ 1 and the admissible length-S pair ξ∈ξ were arbitrary, and the right-hand side is independent of S and ξ, the bound holds uniformly over all such lengths and admissible pairs for the fixed full parameter tuple W, proving Theorem C.3. ∎ C.2 Auxiliary Lemmas for Multi-Layer SSM Federated Convergence First, we describe the following setup that will be utilized in the remainder of the proofs of the multi-layer SSM federated convergence. Fix a reference full-model parameter tuple Wr=(Wr,(1),…,Wr,(Q))W^r=(W^r,(1),…,W^r,(Q)). For each layer q∈1,…,Qq∈\1,…,Q\ of the Q-layer deep SSM defined in C.1, choose some radius Rq>0R_q>0 and define: Rq(q)(Wr,(q)):=V(q):‖V(q)−Wr,(q)‖F≤RqN_R_q^(q) (W^r,(q) ):= \V^(q): \|V^(q)-W^r,(q) \|_F≤ R_q \ (311) as the layerwise parameter neighborhood of radius RqR_q, centered at the parameter tuple of layer q in the central model WrW^r. Furthermore, define: Q _R_Q :=R1(1)(Wr,(1))×⋯×RQ(Q)(Wr,(Q)) :=N_R_1^(1) (W^r,(1) )×…×N_R_Q^(Q) (W^r,(Q) ) (312) Q _Q :=(R1,…,RQ). := (R_1,…,R_Q ). (313) The radii R1,…,RQR_1,…,R_Q are chosen such that: Q⊆,(Q).N_R_Q _ γ, δ^(Q). (314) Furthermore, we assume that Assumption C.1 holds uniformly over this neighborhood with the same finite constants Lx,qL_x,q. Furthermore, we let Assumption 4.2 hold at the final-layer predictions yt(Q)(W,x)y_t^(Q)(W;x) with the same finite constants LℓL_ and HℓH_ throughout this neighborhood. All stability margins, layer-input bounds, and loss-regularity constants are fixed independently of the parameter tuple W∈QW _R_Q, the admissible sequence pair, its finite length, and the token position. In particular, ρq=exp⁡(−γqδq)<1 _q= (- _q _q)<1 is fixed for each layer. For each q∈1,…,Qq∈\1,…,Q\, define the parameter envelopes as: MX,q M_X,q :=‖WXr,(q)‖2+RqX∈Δ,B,C, := \|W_X^r,(q) \|_2+R_q X∈\ ,B,C\, (315) MA,q M_A,q :=‖Ar,(q)‖2+Rq. := \|A^r,(q) \|_2+R_q. (316) By the definition of the layerwise neighborhoods, and the triangle inequality, every W∈QW _R_Q satisfies: ‖WX(q)‖2 \|W_X^(q) \|_2 ≤MX,qX∈Δ,B,C, ≤ M_X,q X∈\ ,B,C\, ‖A(q)‖2 \|A^(q) \|_2 ≤MA,q ≤ M_A,q (317) Given these envelopes, let: MKΔ,q M_K_ ,q :=MΔ,qLx,q+Dlog⁡(2), :=M_ ,qL_x,q+ D (2), (318) MKB,q M_K_B,q :=MB,qLx,q2, :=M_B,qL_x,q^2, (319) MKW,q M_K_W,q :=NDMKB,qMKΔ,q(MKΔ,q+DMA,qLx,q)(ρq1−ρq) :=NDM_K_B,qM_K_ ,q (M_K_ ,q+ DM_A,qL_x,q ) ( _q1- _q ) +Lx,q2NMKΔ,q+Lx,qDMKB,q. +L_x,q^2 NM_K_ ,q+L_x,q DM_K_B,q. (320) These are the layerwise counterparts of the single-layer envelopes introduced in Appendix B.1, with the neighborhood radius R replaced by RqR_q. Then, by the definitions of KΔ,qK_ ,q, KB,qK_B,q, and KW,qK_W,q, it follows that: KΔ,q≤MKΔ,q,KB,q≤MKB,q,KW,q≤MKW,qK_ ,q≤ M_K_ ,q, K_B,q≤ M_K_B,q, K_W,q≤ M_K_W,q (321) over W∈QW _R_Q. Further, define: HRq H_R_q :=MKB,qMKΔ,qND1−ρq, := M_K_B,qM_K_ ,q ND1- _q, (322) JRq J_R_q :=MKW,q1−ρq. := M_K_W,q1- _q. (323) MΛx,q M_ _x,q :=ρqMA,qMΔ,qHRq+MB,qLx,q(2MKΔ,q+Lx,qMΔ,q), := _qM_A,qM_ ,qH_R_q+M_B,qL_x,q (2M_K_ ,q+L_x,qM_ ,q ), (324) Mκ,q M_κ,q :=MC,q(HRq+Lx,qMΛx,q1−ρq). :=M_C,q (H_R_q+ L_x,qM_ _x,q1- _q ). (325) By the definitions of Λx,q _x,q and κq _q in the multi-layer gradient analysis, the following bounds hold Λx,q(W)≤MΛx,q,κq(W)≤Mκ,q _x,q(W)≤ M_ _x,q, _q(W)≤ M_κ,q (326) for every W∈QW _R_Q. These bounds follow by replacing each parameter norm by its corresponding envelope, since all factors are nonnegative. C.2.1 Local multi-layer SSM smoothness Lemma And Proof Lemma C.4 (Local multi-layer SSM smoothness). Under the neighborhood assumptions and parameter envelopes introduced above, there exists a finite constant LSSM,Q(Q)L_SSM,R_Q^(Q) defined: LSSM,Q(Q):=HℓΘQ2+LℓΩQ,L_SSM,R_Q^(Q):=H_ _Q^2+L_ _Q, (327) with ΘQ _Q and ΩQ _Q defined in (330) and (331) respectively, such that, for every admissible sequence pair ξ of finite length S and every W,V∈QW,V _R_Q, ‖∇WℒS,(Q)(W,ξ)−∇WℒS,(Q)(V,ξ)‖F≤LSSM,Q(Q)‖W−V‖F, \| _WL^S,(Q)(W;ξ)- _WL^S,(Q)(V;ξ) \|_F≤ L_SSM,R_Q^(Q)\|W-V\|_F, (328) independent of S and ξ. Furthermore, if each client objective ℒkL_k is finite and differentiable on QN_R_Q and satisfies: ∇Wℒk(W)=ξ∼k[∇WℒS,(Q)(W,ξ)], _WL_k(W)=E_ξ _k [ _WL^S,(Q)(W;ξ) ], (329) then every ℒkL_k, and therefore f=∑kpkℒkf= _kp_kL_k, is also LSSM,Q(Q)L_SSM,R_Q^(Q)-smooth on QN_R_Q. ΘQ _Q and ΩQ _Q have the following recursive definitions: Θ0 _0 :=0,Θq:=(Lx,q2(HRq+MC,qJRq)2+Mκ,q2Θq−12)1/2, :=0, _q:= (L_x,q^2 (H_R_q+M_C,qJ_R_q )^2+M_κ,q^2 _q-1^2 )^1/2, (330) Ω0 _0 :=0,Ωq:=MC,qLx,qℰq1−ρq+2(Lx,q+MC,qΘq−1)(JRq+MΛx,qΘq−11−ρq) :=0, _q:= M_C,qL_x,qE_q1- _q+2 (L_x,q+M_C,q _q-1 ) (J_R_q+ M_ _x,q _q-11- _q ) +2HRqΘq−1+Mκ,qΩq−1. ~~~~~~~~~~~~~~~~~~~~~ +2H_R_q _q-1+M_κ,q _q-1. (331) where: ℰq:= _q:= ρqHRq[(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))2+2(Lx,q+MΔ,qΘq−1) _qH_R_q [ (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) )^2+2 (L_x,q+M_ ,q _q-1 ) +MA,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)] +M_A,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 ) ] +2ρq(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))(JRq+MΛx,qΘq−11−ρq) +2 _q (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) ) (J_R_q+ M_ _x,q _q-11- _q ) +2Θq−1MKΔ,qLx,q +2 _q-1M_K_ ,qL_x,q +2(Lx,q+MB,qΘq−1)[Lx,q(Lx,q+MΔ,qΘq−1)+MKΔ,qΘq−1] +2 (L_x,q+M_B,q _q-1 ) [L_x,q (L_x,q+M_ ,q _q-1 )+M_K_ ,q _q-1 ] +MB,qLx,q[Lx,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)+2(Lx,q+MΔ,qΘq−1)Θq−1]. +M_B,qL_x,q [L_x,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 )+2 (L_x,q+M_ ,q _q-1 ) _q-1 ]. (332) Proof of Lemma C.4: Let W∈QW _R_Q be some fixed arbitrary parameter tuple, and ξ=(x,y∗)ξ=(x,y^*) be some admissible fixed input-target sequence pair of length S≥1S≥ 1. Note that, while x and y∗y^* remain fixed throughout the differentiation, the intermediate inputs and outputs xt(q)=yt(q−1)x_t^(q)=y_t^(q-1) depend on the parameters of the preceding layers for q≥2q≥ 2. We introduce: U=(U(1),…,U(Q)),V=(V(1),…,V(Q))U= (U^(1),…,U^(Q) ), V= (V^(1),…,V^(Q) ) (333) as two arbitrary full-model perturbations such that: U(q)=(UΔ(q),UB(q),UC(q),UA(q)),V(q)=(VΔ(q),VB(q),VC(q),VA(q))U^(q)= (U_ ^(q),U_B^(q),U_C^(q),U_A^(q) ), V^(q)= (V_ ^(q),V_B^(q),V_C^(q),V_A^(q) ) (334) and: ‖U‖F2:=∑q=1Q‖U(q)‖F2=1,‖V‖F2:=∑q=1Q‖V(q)‖F2=1.\|U\|_F^2:= _q=1^Q \|U^(q) \|_F^2=1, \|V\|_F^2:= _q=1^Q \|V^(q) \|_F^2=1. (335) Note that each layerwise norm such as ‖U(q)‖F \|U^(q) \|_F is defined as the Euclidean combination of the Frobenius norms of its four parameter blocks. As a result, every layerwise direction and every individual parameter-block direction has norm at most one. Then, for a twice continuously differentiable function F of the full parameter tuple, we define the first and second directional derivatives by DF(W)[U] DF(W)[U] :=dϵF(W+ϵU)|ϵ=0, := ddεF(W+ε U) |_ε=0, (336) D2F(W)[U,V] D^2F(W)[U,V] :=∂2∂ϵ∂ζF(W+ϵU+ζV)|ϵ=ζ=0. := ∂^2∂ε∂ζF(W+ε U+ζ V) |_ε=ζ=0. (337) Unless otherwise specified, all derivatives in the remaining part of this proof are taken with respect to the full parameter tuple W. The first-order chain rule gives D⁡(ℓt(yt(Q)))(W)[U]=⟨∇yℓt(yt(Q)(W)),Dyt(Q)(W)[U]⟩.D ( _t (y_t^(Q) ) )(W)[U]= _y _t (y_t^(Q)(W) ),Dy_t^(Q)(W)[U] . (338) where we used the notation ℓt(y)=ℓt(y,yt∗) _t(y)= _t(y,y_t^*). By differentiating this expression in V direction, we get: D2(ℓt(yt(Q)))(W)[U,V]= D^2 ( _t (y_t^(Q) ) )(W)[U,V]= (Dyt(Q)(W)[U])⊤∇y2ℓt(yt(Q)(W))Dyt(Q)(W)[V] (Dy_t^(Q)(W)[U] ) _y^2 _t (y_t^(Q)(W) )Dy_t^(Q)(W)[V] +⟨∇yℓt(yt(Q)(W)),D2yt(Q)(W)[U,V]⟩. + _y _t (y_t^(Q)(W) ),D^2y_t^(Q)(W)[U,V] . (339) By the neighborhood-uniform loss-regularity bounds introduced in Assumption 4.2 and the Cauchy-Schwarz inequality, it follows that |D2(ℓt(yt(Q)))(W)[U,V]|≤ |D^2 ( _t (y_t^(Q) ) )(W)[U,V] |≤ Hℓ‖Dyt(Q)(W)[U]‖2‖Dyt(Q)(W)[V]‖2+Lℓ‖D2yt(Q)(W)[U,V]‖2. H_ \|Dy_t^(Q)(W)[U] \|_2 \|Dy_t^(Q)(W)[V] \|_2+L_ \|D^2y_t^(Q)(W)[U,V] \|_2. (340) As such, to bound the second derivative of the loss, it suffices to derive uniform bounds on the first and second directional derivatives of the final-layer output. We first establish the first-order output sensitivity bounds needed for the second-order calculation. For ease of notation, for a full-model direction U we write U(≤q)=(U(1),…,U(q))U^(≤ q)= (U^(1),…,U^(q) ) and U(<q)=(U(1),…,U(q−1))U^(<q)=(U^(1),…,U^(q-1)). We further define the following recursion: Θ0:=0,Θq:=(Lx,q2(HRq+MC,qJRq)2+Mκ,q2Θq−12)1/2 _0:=0, _q:= (L_x,q^2 (H_R_q+M_C,qJ_R_q )^2+M_κ,q^2 _q-1^2 )^1/2 (341) First, by (270) and the definition of the neighborhood envelopes, we have: ‖∂yt(q)∂W(q)‖F≤Lx,q(HRq+MC,qJRq). \| ∂ y_t^(q)∂ W^(q) \|_F≤ L_x,q (H_R_q+M_C,qJ_R_q ). (342) Then, by (304) and the bound κq(W)≤Mκ,q _q(W)≤ M_κ,q, for every t∈1,…,St∈\1,…,S\, ∑j=1t‖∂yt(q)∂xj(q)‖2=∑j=1t‖t,j(q)‖2≤κq(W)≤Mκ,q. _j=1^t \| ∂ y_t^(q)∂ x_j^(q) \|_2= _j=1^t \|K_t,j^(q) \|_2≤ _q(W)≤ M_κ,q. (343) We next prove by induction on q that ‖Dyt(q)(W)[U]‖2≤Θq‖U(≤q)‖F \|Dy_t^(q)(W)[U] \|_2≤ _q \|U^(≤ q) \|_F (344) for every q∈1,…,Qq∈\1,…,Q\ and every t∈1,…,St∈\1,…,S\. When q=1q=1, the input sequence is independent of W, and the output of the first layer depends only on W(1)W^(1). As such: Dyt(1)(W)[U]=∂yt(1)∂W(1)vec⁡(U(1)),Dy_t^(1)(W)[U]= ∂ y_t^(1)∂ W^(1)vec (U^(1) ), (345) Using ‖Az‖2≤‖A‖F‖z‖2\|Az\|_2≤\|A\|_F\|z\|_2 and ‖vec⁡(U(1))‖2=‖U(1)‖F\|vec(U^(1))\|_2=\|U^(1)\|_F gives ‖Dyt(1)(W)[U]‖2≤Lx,1(HR1+MC,1JR1)‖U(1)‖F=Θ1‖U(≤1)‖F, \|Dy_t^(1)(W)[U] \|_2≤ L_x,1 (H_R_1+M_C,1J_R_1 ) \|U^(1) \|_F= _1 \|U^(≤ 1) \|_F, (346) where the equality follows from Θ0=0 _0=0 and the definition of Θ1 _1. Next, for a fixed q≥2q≥ 2, Since U(≤q−1)=U(<q)U^(≤ q-1)=U^(<q), suppose (the induction hypothesis): ‖Dyj(q−1)(W)[U]‖2≤Θq−1‖U(<q)‖F,j∈1,…,S. \|Dy_j^(q-1)(W)[U] \|_2≤ _q-1 \|U^(<q) \|_F, j∈\1,…,S\. (347) For a fixed t, the output yt(q)y_t^(q) depends on the current-layer parameters W(q)W^(q) and the input tokens x1(q),…,xt(q)x_1^(q),…,x_t^(q), which depend on the preceding layers through xj(q)=yj(q−1)x_j^(q)=y_j^(q-1). The chain rule therefore gives Dyt(q)(W)[U]=∂yt(q)∂W(q)vec⁡(U(q))+∑j=1t∂yt(q)∂xj(q)Dyj(q−1)(W)[U].Dy_t^(q)(W)[U]= ∂ y_t^(q)∂ W^(q)vec (U^(q) )+ _j=1^t ∂ y_t^(q)∂ x_j^(q)Dy_j^(q-1)(W)[U]. (348) The first term accounts for the direct parameter perturbation in layer q, while the sum accounts for changes propagated through its input sequence. Taking norms and applying the established Jacobian bounds and the induction hypothesis yields ‖Dyt(q)(W)[U]‖2≤ \|Dy_t^(q)(W)[U] \|_2≤ Lx,q(HRq+MC,qJRq)‖U(q)‖F+∑j=1t‖∂yt(q)∂xj(q)‖2‖Dyj(q−1)(W)[U]‖2 L_x,q (H_R_q+M_C,qJ_R_q ) \|U^(q) \|_F+ _j=1^t \| ∂ y_t^(q)∂ x_j^(q) \|_2 \|Dy_j^(q-1)(W)[U] \|_2 ≤ ≤ Lx,q(HRq+MC,qJRq)‖U(q)‖F+Mκ,qΘq−1‖U(<q)‖F. L_x,q (H_R_q+M_C,qJ_R_q ) \|U^(q) \|_F+M_κ,q _q-1 \|U^(<q) \|_F. (349) Applying Cauchy-Schwarz to the two terms gives ‖Dyt(q)(W)[U]‖2 \|Dy_t^(q)(W)[U] \|_2 ≤(Lx,q2(HRq+MC,qJRq)2+Mκ,q2Θq−12)1/2×(‖U(q)‖F2+‖U(<q)‖F2)1/2 ≤ (L_x,q^2 (H_R_q+M_C,qJ_R_q )^2+M_κ,q^2 _q-1^2 )^1/2× ( \|U^(q) \|_F^2+ \|U^(<q) \|_F^2 )^1/2 =Θq‖U(≤q)‖F, = _q \|U^(≤ q) \|_F, (350) where the equality follows from the definition of Θq _q and the block Frobenius norm. This completes the induction. The same argument applies to V. Finally, since ‖U(≤q)‖F≤‖U‖F=1\|U^(≤ q)\|_F≤\|U\|_F=1 and ‖V(≤q)‖F≤‖V‖F=1\|V^(≤ q)\|_F≤\|V\|_F=1, we obtain the following first-derivative bounds: ‖Dyt(q)(W)[U]‖2≤Θq,‖Dyt(q)(W)[V]‖2≤Θq, \|Dy_t^(q)(W)[U] \|_2≤ _q, \|Dy_t^(q)(W)[V] \|_2≤ _q, (351) uniformly over all layers q∈1,…,Qq∈\1,…,Q\ and token positions t∈1,…,St∈\1,…,S\. Finally, applying this result to q=Qq=Q, we get the resulting first-derivative bounds: ‖Dyt(Q)(W)[U]‖2≤ΘQ,‖Dyt(Q)(W)[V]‖2≤ΘQ, \|Dy_t^(Q)(W)[U] \|_2≤ _Q, \|Dy_t^(Q)(W)[V] \|_2≤ _Q, (352) and substituting back into (340) we get: |D2(ℓt(yt(Q)))(W)[U,V]|≤ |D^2 ( _t (y_t^(Q) ) )(W)[U,V] |≤ HℓΘQ2+Lℓ‖D2yt(Q)(W)[U,V]‖2. H_ _Q^2+L_ \|D^2y_t^(Q)(W)[U,V] \|_2. (353) We next bound the second output derivative term ‖D2yt(Q)(W)[U,V]‖2 \|D^2y_t^(Q)(W)[U,V] \|_2. For any layer q, the second-order product rule gives: D2yt(q)(W)[U,V]= D^2y_t^(q)(W)[U,V]= (D2ht(q)(W)[U,V])⊤Ct(q)+(ht(q))⊤D2Ct(q)(W)[U,V] (D^2h_t^(q)(W)[U,V] ) C_t^(q)+ (h_t^(q) ) D^2C_t^(q)(W)[U,V] +(Dht(q)(W)[U])⊤DCt(q)(W)[V]+(Dht(q)(W)[V])⊤DCt(q)(W)[U]. + (Dh_t^(q)(W)[U] ) DC_t^(q)(W)[V]+ (Dh_t^(q)(W)[V] ) DC_t^(q)(W)[U]. (354) We bound each term separately. By the first-order bounds already established: ‖Dxt(q)(W)[U]‖2=‖Dyt(q−1)(W)[U]‖2≤Θq−1, \|Dx_t^(q)(W)[U] \|_2= \|Dy_t^(q-1)(W)[U] \|_2≤ _q-1, (355) vanishing for q=1q=1, consistent with Θ0=0 _0=0. The same holds for ‖Dxt(q)(W)[V]‖2\|Dx_t^(q)(W)[V]\|_2. As a result, by the product rule: ‖DCt(q)(W)[U]‖2 \|DC_t^(q)(W)[U] \|_2 ≤Lx,q+MC,qΘq−1, ≤ L_x,q+M_C,q _q-1, (356) ‖DCt(q)(W)[V]‖2 \|DC_t^(q)(W)[V] \|_2 ≤Lx,q+MC,qΘq−1, ≤ L_x,q+M_C,q _q-1, (357) ‖D2Ct(q)(W)[U,V]‖2 \|D^2C_t^(q)(W)[U,V] \|_2 ≤2Θq−1+MC,q‖D2xt(q)(W)[U,V]‖2. ≤ 2 _q-1+M_C,q \|D^2x_t^(q)(W)[U,V] \|_2. (358) Furthermore, by (268) and (297), ‖Dht(q)(W)[U]‖F \|Dh_t^(q)(W)[U] \|_F ≤JRq+MΛx,qΘq−1∑j=1tρqt−j≤JRq+MΛx,qΘq−11−ρq, ≤ J_R_q+M_ _x,q _q-1 _j=1^t _q^t-j≤ J_R_q+ M_ _x,q _q-11- _q, (359) with the same inequality holding for ‖Dht(q)(W)[V]‖F \|Dh_t^(q)(W)[V] \|_F. Since ‖Ct(q)‖2≤MC,qLx,q \|C_t^(q) \|_2≤ M_C,qL_x,q and ‖ht(q)‖F≤HRq \|h_t^(q) \|_F≤ H_R_q, equation (354) yields ‖D2yt(q)(W)[U,V]‖2≤ \|D^2y_t^(q)(W)[U,V] \|_2≤ MC,qLx,q‖D2ht(q)(W)[U,V]‖F+MC,qHRq‖D2xt(q)(W)[U,V]‖2 M_C,qL_x,q \|D^2h_t^(q)(W)[U,V] \|_F+M_C,qH_R_q \|D^2x_t^(q)(W)[U,V] \|_2 +2(Lx,q+MC,qΘq−1)(JRq+MΛx,qΘq−11−ρq)+2HRqΘq−1. +2 (L_x,q+M_C,q _q-1 ) (J_R_q+ M_ _x,q _q-11- _q )+2H_R_q _q-1. (360) As a result, we need to bound the second derivative of the hidden state. By the definition of ht(q)(W)h_t^(q)(W), differentiating it twice yields: D2ht(q)(W)[U,V]= D^2h_t^(q)(W)[U,V]= A¯t(q)⊙D2ht−1(q)(W)[U,V]+D2A¯t(q)(W)[U,V]⊙ht−1(q) A_t^(q) D^2h_t-1^(q)(W)[U,V]+D^2 A_t^(q)(W)[U,V] h_t-1^(q) +DA¯t(q)(W)[U]⊙Dht−1(q)(W)[V]+DA¯t(q)(W)[V]⊙Dht−1(q)(W)[U] +D A_t^(q)(W)[U] Dh_t-1^(q)(W)[V]+D A_t^(q)(W)[V] Dh_t-1^(q)(W)[U] +D2[Bt(q)(Δt(q)⊙xt(q))⊤](W)[U,V]. +D^2 [B_t^(q) ( _t^(q) x_t^(q) ) ](W)[U,V]. (361) Then, since the absolute values of the first and second derivatives of SoftplusSoftplus are bounded above by 11 and 1/41/4, respectively, the product and chain rule give: ‖DΔt(q)(W)[U]‖2 \|D _t^(q)(W)[U] \|_2 ≤Lx,q+MΔ,qΘq−1, ≤ L_x,q+M_ ,q _q-1, (362) ‖DΔt(q)(W)[V]‖2 \|D _t^(q)(W)[V] \|_2 ≤Lx,q+MΔ,qΘq−1, ≤ L_x,q+M_ ,q _q-1, (363) ‖D2Δt(q)(W)[U,V]‖2 \|D^2 _t^(q)(W)[U,V] \|_2 ≤14(Lx,q+MΔ,qΘq−1)2+2Θq−1+MΔ,q‖D2xt(q)(W)[U,V]‖2. ≤ 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1+M_ ,q \|D^2x_t^(q)(W)[U,V] \|_2. (364) Then, applying the product rule and the exponential chain rule (per every entry of A), and using the established state and first-derivative bounds, we get: ‖D2ht(q)(W)[U,V]‖F≤ \|D^2h_t^(q)(W)[U,V] \|_F≤ ρq‖D2ht−1(q)(W)[U,V]‖F+ℰq+MΛx,q‖D2xt(q)(W)[U,V]‖2, _q \|D^2h_t-1^(q)(W)[U,V] \|_F+E_q+M_ _x,q \|D^2x_t^(q)(W)[U,V] \|_2, (365) where we collected all the terms independent of the second input derivative in ℰqE_q defined as ℰq:= _q:= ρqHRq[(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))2+2(Lx,q+MΔ,qΘq−1) _qH_R_q [ (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) )^2+2 (L_x,q+M_ ,q _q-1 ) +MA,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)] +M_A,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 ) ] +2ρq(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))(JRq+MΛx,qΘq−11−ρq) +2 _q (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) ) (J_R_q+ M_ _x,q _q-11- _q ) +2Θq−1MKΔ,qLx,q +2 _q-1M_K_ ,qL_x,q +2(Lx,q+MB,qΘq−1)[Lx,q(Lx,q+MΔ,qΘq−1)+MKΔ,qΘq−1] +2 (L_x,q+M_B,q _q-1 ) [L_x,q (L_x,q+M_ ,q _q-1 )+M_K_ ,q _q-1 ] +MB,qLx,q[Lx,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)+2(Lx,q+MΔ,qΘq−1)Θq−1]. +M_B,qL_x,q [L_x,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 )+2 (L_x,q+M_ ,q _q-1 ) _q-1 ]. (366) The first term in ℰqE_q comes from the second transition derivative multiplied entrywise by the state: D2A¯t(q)(W)[U,V]⊙ht−1(q).D^2 A_t^(q)(W)[U,V] h_t-1^(q). The second comes from the two mixed transition-state terms, DA¯t(q)(W)[U]⊙Dht−1(q)(W)[V]+DA¯t(q)(W)[V]⊙Dht−1(q)(W)[U].D A_t^(q)(W)[U] Dh_t-1^(q)(W)[V]+D A_t^(q)(W)[V] Dh_t-1^(q)(W)[U]. The remaining three terms come from differentiating the additive state term Bt(q)(Δt(q)⊙xt(q))⊤B_t^(q) ( _t^(q) x_t^(q) ) twice. The contributions involving the second input derivative are collectively bounded by MΛx,q‖D2xt(q)(W)[U,V]‖2,M_ _x,q \|D^2x_t^(q)(W)[U,V] \|_2, where MΛx,qM_ _x,q was defined in (324). Since the initial state and its derivatives equal 0, unrolling (365) and summing the geometric series gives ‖D2ht(q)(W)[U,V]‖F≤ℰq1−ρq+MΛx,q1−ρqmax1≤j≤t‖D2xj(q)(W)[U,V]‖2. \|D^2h_t^(q)(W)[U,V] \|_F≤ E_q1- _q+ M_ _x,q1- _q _1≤ j≤ t \|D^2x_j^(q)(W)[U,V] \|_2. (367) Next, we note that since ‖D2xt(q)(W)[U,V]‖2≤max1≤j≤t⁡‖D2xj(q)(W)[U,V]‖2 \|D^2x_t^(q)(W)[U,V] \|_2≤ _1≤ j≤ t \|D^2x_j^(q)(W)[U,V] \|_2, substituting (367) and using this fact, the new coefficient of the max1≤j≤t⁡‖D2xj(q)(W)[U,V]‖2 _1≤ j≤ t \|D^2x_j^(q)(W)[U,V] \|_2 term becomes: MC,qHRq+MC,qLx,qMΛx,q1−ρq=Mκ,q.M_C,qH_R_q+ M_C,qL_x,qM_ _x,q1- _q=M_κ,q. (368) Finally, we define a new recursive term Ωq _q as follows: Ω0 _0 :=0, :=0, Ωq _q :=MC,qLx,qℰq1−ρq+2(Lx,q+MC,qΘq−1)(JRq+MΛx,qΘq−11−ρq) := M_C,qL_x,qE_q1- _q+2 (L_x,q+M_C,q _q-1 ) (J_R_q+ M_ _x,q _q-11- _q ) +2HRqΘq−1+Mκ,qΩq−1. +2H_R_q _q-1+M_κ,q _q-1. (369) For q=1q=1, second input derivative vanishes. Then, for q≥2q≥ 2, and using the similar induction argument as in the first-order derivative, as well as the fact xj(q)=yj(q−1)x_j^(q)=y_j^(q-1), the induction yields: ‖D2yt(q)(W)[U,V]‖2≤Ωq \|D^2y_t^(q)(W)[U,V] \|_2≤ _q (370) for every layer and token position. All constants are finite and independent of W, ξ, S, and t, because the neighborhood envelopes are fixed, ρq<1 _q<1, and Q is finite. Finally, applying this bound to q=Qq=Q, and substituting the result into (340) it follows that: |D2(ℓt(yt(Q)))(W)[U,V]|≤HℓΘQ2+LℓΩQ=:LSSM,Q(Q). |D^2 ( _t (y_t^(Q) ) )(W)[U,V] |≤ H_ _Q^2+L_ _Q=:L_SSM,R_Q^(Q). (371) Averaging over the token losses preserves this bound. Since U,VU,V are arbitrary unit directions, it is true that: ‖∇W2ℒS,(Q)(W,ξ)‖op≤LSSM,Q(Q), \| _W^2L^S,(Q)(W;ξ) \|_op≤ L_SSM,R_Q^(Q), (372) where the operator norm is induced by the block Frobenius norm. Since QN_R_Q is convex, integrating the Hessian along the segment joining any W0,W1∈QW_0,W_1 _R_Q gives ‖∇WℒS,(Q)(W1,ξ)−∇WℒS,(Q)(W0,ξ)‖F≤LSSM,Q(Q)‖W1−W0‖F. \| _WL^S,(Q)(W_1;ξ)- _WL^S,(Q)(W_0;ξ) \|_F≤ L_SSM,R_Q^(Q)\|W_1-W_0\|_F. (373) Taking expectations, by Assumption 4.6 and Jensen’s inequality, we get: ‖∇Wℒk(W1)−∇Wℒk(W0)‖F \| _WL_k(W_1)- _WL_k(W_0) \|_F =‖ξ∼k[∇WℒS,(Q)(W1,ξ)−∇WℒS,(Q)(W0,ξ)]‖F = \|E_ξ _k [ _WL^S,(Q)(W_1;ξ)- _WL^S,(Q)(W_0;ξ) ] \|_F ≤ξ∼k[‖∇WℒS,(Q)(W1,ξ)−∇WℒS,(Q)(W0,ξ)‖F] _ξ _k [ \| _WL^S,(Q)(W_1;ξ)- _WL^S,(Q)(W_0;ξ) \|_F ] ≤LSSM,Q(Q)‖W1−W0‖F, ≤ L_SSM,R_Q^(Q)\|W_1-W_0\|_F, (374) demonstrating that every client objective is LSSM,Q(Q)L_SSM,R_Q^(Q)-smooth on QN_R_Q. Finally, since this inequality holds for every client objective, and pk≥0p_k≥ 0 and ∑kpk=1 _kp_k=1, it also holds for f=∑kpkℒkf= _kp_kL_k, demonstrating that f is also LSSM,Q(Q)L_SSM,R_Q^(Q)-smooth on QN_R_Q, completing the proof. ∎. C.2.2 Multi-layer SSM gradient second-moment constant Lemma and Proof Lemma C.5 (Multi-layer SSM gradient second-moment constant). Suppose the Assumptions 4.2 and C.1 hold. Let Q(Wr)⊆,(Q)N_R_Q(W^r) ^(Q)_ γ, δ be a stable local model neighborhood around the round-r broadcast Q-deep model WrW^r. Let: g⁡(W,ξ):=∇WℒS,(Q)(W,ξ)g(W;ξ):= _WL^S,(Q)(W;ξ) (375) denote the stochastic sequence-level gradient for a sampled input-target sequence pair ξ=(x1:S,y1:S∗)ξ=(x_1:S,y_1:S^*). Then, for every W∈Q(Wr)W _R_Q(W^r), ξ‖g⁡(W,ξ)‖F2≤(GSSM,Q(Q))2,E_ξ \|g(W;ξ) \|_F^2≤ (G_SSM,R_Q^(Q) )^2, (376) where GSSM,Q(Q):=Lℓ(∑q=1QLx,q2(HRq+MC,qJRq)2∏q′=q+1QMκ,q′2)1/2=LℓΘQ.G_SSM,R_Q^(Q):=L_ ( _q=1^QL_x,q^2 (H_R_q+M_C,qJ_R_q )^2 _q =q+1^QM_κ,q ^2 )^1/2=L_ _Q. (377) where HRq,JRq,MC,q and Mκ,qH_R_q,J_R_q,M_C,q and M_κ,q are defined in (322), (323), (315), and (325) respectively. Proof of Lemma C.5: By Theorem C.3, under the assumptions of the lemma, for every W∈Q(Wr)⊆,(Q)W _R_Q(W^r) ^(Q)_ γ, δ, every finite integer S≥1S≥ 1, and every admissible length-S sequence pair ξ∈ξ , we have: ‖∇WℒS,(Q)(W,ξ)‖F≤Lℓ(∑q=1QLx,q2(Hq+‖WC(q)‖2Jq)2∏q′=q+1Qκq′2)1/2. \| _WL^S,(Q)(W;ξ) \|_F≤ L_ ( _q=1^QL_x,q^2 (H_q+ \|W_C^(q) \|_2J_q )^2 _q =q+1^Q _q ^2 )^1/2. Since g⁡(W,ξ)=∇WℒS,(Q)(W,ξ)g(W;ξ)= _WL^S,(Q)(W;ξ), it follows that: ‖g⁡(W,ξ)‖F≤Lℓ(∑q=1QLx,q2(Hq+‖WC(q)‖2Jq)2∏q′=q+1Qκq′2)1/2.\|g(W;ξ)\|_F≤ L_ ( _q=1^QL_x,q^2 (H_q+ \|W_C^(q) \|_2J_q )^2 _q =q+1^Q _q ^2 )^1/2. (378) Restricting to W∈Q(Wr)W _R_Q(W^r), by the local neighborhood bounds defined previously ( (322), (323), (315), and (325)), Hq≤HRq,‖WC(q)‖2≤MC,q,Jq≤JRq,κq≤Mκ,q. H_q≤ H_R_q, \|W_C^(q) \|_2≤ M_C,q, J_q≤ J_R_q, _q≤ M_κ,q. (379) Therefore, ‖g⁡(W,ξ)‖F≤Lℓ(∑q=1QLx,q2(HRq+MC,qJRq)2∏q′=q+1QMκ,q′2)1/2=GSSM,Q(Q). \|g(W;ξ) \|_F≤ L_ ( _q=1^QL_x,q^2 (H_R_q+M_C,qJ_R_q )^2 _q =q+1^QM_κ,q ^2 )^1/2=G_SSM,R_Q^(Q). (380) Since both values are non-negative, squaring both sides gives: ‖g⁡(W,ξ)‖F2≤(GSSM,Q(Q))2=Lℓ2(∑q=1QLx,q2(HRq+MC,qJRq)2∏q′=q+1QMκ,q′2)\|g(W;ξ)\|_F^2≤ (G_SSM,R_Q^(Q) )^2=L_ ^2 ( _q=1^QL_x,q^2 (H_R_q+M_C,qJ_R_q )^2 _q =q+1^QM_κ,q ^2 ) (381) Finally, taking expectation over the sampled sequence ξ preserves the bound: ξ[‖g⁡(W,ξ)‖F2]≤(GSSM,Q(Q))2E_ξ [\|g(W;ξ)\|_F^2 ]≤ (G_SSM,R_Q^(Q) )^2 (382) proving the lemma. ∎ C.2.3 Multi-layer SSM stochastic gradient variance bound Lemma and Proof Lemma C.6 (Multi-layer SSM stochastic gradient variance bound). Suppose the assumptions of Lemma C.5 and Assumption 4.6 hold. Let ℒ⁡(W):=ξ∼[ℒS,(Q)(W,ξ)]L(W):=E_ξ [L^S,(Q)(W;ξ) ] denote the expected sequence loss over some local data distribution P, for a Q-layer deep SSM parameterized by W. Define g⁡(W,ξ):=∇WℒS,(Q)(W,ξ)g(W;ξ):= _WL^S,(Q)(W;ξ) as the single-sequence stochastic gradient. For a mini-batch Ξξ=ξ1,…,ξΞ _ξ=\ _1,…, _ \ of Ξ independent sampled sequences, define gΞ(W,Ξξ):=1Ξ∑j=1Ξg⁡(W,ξj).g_ (W; _ξ):= 1 _j=1 g(W; _j). (383) Then, for every W∈Q(Wr)W _R_Q(W^r) it follows that Ξξ[‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2]≤(σSSM,Q,Ξ(Q))2,E_ _ξ [ \|g_ (W; _ξ)- _WL(W) \|_F^2 ]≤ ( _SSM,R_Q, ^(Q) )^2, (384) where (σSSM,Q,Ξ(Q))2:=1ΞsupW∈Q(Wr)(ξ∼‖g⁡(W,ξ)‖F2−‖∇Wℒ(W)‖F2). ( _SSM,R_Q, ^(Q) )^2:= 1 _W _R_Q(W^r) (E_ξ \|g(W;ξ)\|_F^2-\| _WL(W)\|_F^2 ). (385) Moreover, by Lemma C.5, (σSSM,Q,Ξ(Q))2≤(GSSM,Q(Q))2Ξ. ( _SSM,R_Q, ^(Q) )^2≤ (G_SSM,R_Q^(Q) )^2 . (386) Proof of Lemma C.6: By unbiasedness and independence of the sampled sequences, the standard mini-batch variance identity gives: Ξξ[‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2]=1Ξξ∼[‖g⁡(W,ξ)−∇Wℒ(W)‖F2]. _ _ξ [ \|g_ (W; _ξ)- _WL(W) \|_F^2 ]= 1 E_ξ [ \|g(W;ξ)- _WL(W) \|_F^2 ]. Using the identity ξ∼[‖g⁡(W,ξ)−∇Wℒ(W)‖F2]=ξ∼‖g⁡(W,ξ)‖F2−‖∇Wℒ(W)‖F2,E_ξ [ \|g(W;ξ)- _WL(W) \|_F^2 ]=E_ξ \|g(W;ξ)\|_F^2-\| _WL(W)\|_F^2, (387) and taking the supremum over W∈Q(Wr)W _R_Q(W^r), we obtain Ξξ[‖gΞ(W,Ξξ)−∇Wℒ(W)‖F2]≤(σSSM,Q,Ξ(Q))2.E_ _ξ [ \|g_ (W; _ξ)- _WL(W) \|_F^2 ]≤ ( _SSM,R_Q, ^(Q) )^2. (388) Finally, Lemma C.5 gives ξ∼‖g⁡(W,ξ)‖F2≤(GSSM,Q(Q))2E_ξ \|g(W;ξ)\|_F^2≤ (G_SSM,R_Q^(Q) )^2 and since ‖∇Wℒ(W)‖F2≥0\| _WL(W)\|_F^2≥ 0, (σSSM,Q,Ξ(Q))2≤(GSSM,Q(Q))2Ξ. ( _SSM,R_Q, ^(Q) )^2≤ (G_SSM,R_Q^(Q) )^2 . (389) This proves the claim. ∎ C.2.4 Multi-layer SSM client heterogeneity bound under total variation Lemma and Proof Lemma C.7 (Multi-layer SSM client heterogeneity bound under total variation). Suppose the assumptions of Lemma C.5 and Assumption 4.6 hold. Let f⁡(W):=∑k=1Kpkℒk(W)f(W):= _k=1^Kp_kL_k(W) (390) be the global objective, where pk≥0p_k≥ 0 and ∑k=1Kpk=1 _k=1^Kp_k=1. Let kP_k denote the local data distribution of client k, and define the global mixture distribution as: ¯:=∑k=1Kpkk. P:= _k=1^Kp_kP_k. (391) Define the local client heterogeneity constant (ΓSSM,Q(Q))2:=supW∈Q(Wr)∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2. ( _SSM,R_Q^(Q) )^2:= _W _R_Q(W^r) _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2. (392) Then: (ΓSSM,Q(Q))2≤4(GSSM,Q(Q))2∑k=1KpkTV⁡(k,¯)2, ( _SSM,R_Q^(Q) )^2≤ 4 (G_SSM,R_Q^(Q) )^2 _k=1^Kp_kTV(P_k, P)^2, (393) where TV⁡(P,Q):=supA|P⁡(A)−Q⁡(A)|.TV(P,Q):= _A|P(A)-Q(A)|. (394) Proof of Lemma C.7: Recall that for each client k, the local objective is: ℒk(W)=ξ∼k[ℒS,(Q)(W,ξ)],L_k(W)=E_ξ _k [L^S,(Q)(W;ξ) ], and, by Assumption 4.6, ∇Wℒk(W)=ξ∼k[g⁡(W,ξ)]. _WL_k(W)=E_ξ _k [g(W;ξ) ]. (395) Similarly, since f⁡(W):=∑k=1Kpkℒk(W),f(W):= _k=1^Kp_kL_k(W), we have ∇Wf(W)=∑k=1Kpk∇Wℒk(W)=∑k=1Kpkξ∼k[g⁡(W,ξ)]. _Wf(W)= _k=1^Kp_k _WL_k(W)= _k=1^Kp_kE_ξ _k [g(W;ξ) ]. (396) Next, by the definition of the mixture distribution ¯ P, we can write this as: ∇Wf(W)=ξ∼¯[g⁡(W,ξ)]. _Wf(W)=E_ξ P [g(W;ξ) ]. (397) Therefore, it follows that ∇Wℒk(W)−∇Wf(W)=ξ∼k[g⁡(W,ξ)]−ξ∼¯[g⁡(W,ξ)]. _WL_k(W)- _Wf(W)=E_ξ _k[g(W;ξ)]-E_ξ P[g(W;ξ)]. (398) By the pointwise bound (380) established in the proof of Lemma C.5, for every W∈Q(Wr)W _R_Q(W^r) and every admissible sequence pair ξ∈ξ , ‖g⁡(W,ξ)‖F≤GSSM,Q(Q). \|g(W;ξ) \|_F≤ G_SSM,R_Q^(Q). Since k()=¯()=1P_k(Z)= P(Z)=1, the total-variation inequality for bounded vector-valued functions gives: ‖ξ∼k[g⁡(W,ξ)]−ξ∼¯[g⁡(W,ξ)]‖F≤2GSSM,Q(Q)TV⁡(k,¯). \|E_ξ _k[g(W;ξ)]-E_ξ P[g(W;ξ)] \|_F≤ 2G_SSM,R_Q^(Q)TV(P_k, P). (399) So: ‖∇Wℒk(W)−∇Wf(W)‖F2≤4(GSSM,Q(Q))2TV⁡(k,¯)2. \| _WL_k(W)- _Wf(W) \|_F^2≤ 4 (G_SSM,R_Q^(Q) )^2TV(P_k, P)^2. (400) Multiplying both sides by pkp_k and summing over all clients gives: ∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2≤4(GSSM,Q(Q))2∑k=1KpkTV⁡(k,¯)2. _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2≤ 4 (G_SSM,R_Q^(Q) )^2 _k=1^Kp_kTV(P_k, P)^2. (401) Since the right-hand side of this inequality is independent of W, taking the supremum over W∈Q(Wr)W _R_Q(W^r) we get: (ΓSSM,Q(Q))2:=supW∈Q(Wr)∑k=1Kpk‖∇Wℒk(W)−∇Wf(W)‖F2≤4(GSSM,Q(Q))2∑k=1KpkTV⁡(k,¯)2. ( _SSM,R_Q^(Q) )^2:= _W _R_Q(W^r) _k=1^Kp_k \| _WL_k(W)- _Wf(W) \|_F^2≤ 4 (G_SSM,R_Q^(Q) )^2 _k=1^Kp_kTV(P_k, P)^2. (402) This proves the multi-layer SSM client heterogeneity bound. ∎ C.3 Corollary C.8: FedAvg Multi-Layer SSM Convergence and Proof Corollary C.8 (FedAvg convergence for multi-layer selective SSMs). Consider the distributed learning framework established in Section 3, with objective f⁡(W):=∑k=1Kpkℒk(W),pk≥0,∑k=1Kpk=1,f(W):= _k=1^Kp_kL_k(W), p_k≥ 0, _k=1^Kp_k=1, where the weights pkp_k are fixed and W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ) parameterizes the selective SSM of finite depth Q≥1Q≥ 1 defined in Appendix C.1. Each client objective is finite and differentiable on the neighborhoods under consideration, and satisfies ∇Wℒk(W)=ξ∼k[∇WℒS,(Q)(W,ξ)]. _WL_k(W)=E_ξ _k [ _WL^S,(Q)(W;ξ) ]. where S is the length of the sampled sequence pair ξ, and all clients use the same sequence-loss map. Let W0W^0 be deterministic with f⁡(W0)<∞f(W^0)<∞, and suppose that f is bounded below by a finite constant f∗f^*. Consider full-participation FedAvg for T communication rounds, with E local steps per round, mini-batch size Ξ , and constant deterministic local step size η>0η>0, where T,E,Ξ∈ℤ>0T,E, _>0: Wkr,0 W_k^r,0 =Wr, =W^r, Wkr,e+1 W_k^r,e+1 =Wkr,e−ηgk,Ξr,e,e=0,…,E−1, =W_k^r,e-η g_k, ^r,e, e=0,…,E-1, Wr+1 W^r+1 =∑k=1KpkWkr,E. = _k=1^Kp_kW_k^r,E. The mini-batch gradient is gk,Ξr,e:=1Ξ∑i=1Ξ∇WℒSi,(Q)(Wkr,e,ξk,ir,e),g_k, ^r,e:= 1 _i=1 _WL^S_i,(Q)(W_k^r,e; _k,i^r,e), where SiS_i denotes the length of ξk,ir,e _k,i^r,e. Conditional on the complete training history before each local step, each client draws Ξ fresh independent samples from kP_k, and the mini-batches drawn by different clients are conditionally independent. For fixed deterministic radii Rq>0R_q>0, define Q(Wr):=∏q=1QRq(q)(Wr,(q)),Q:=(R1,…,RQ). _R_Q(W^r):= _q=1^QN_R_q^(q)(W^r,(q)), _Q:=(R_1,…,R_Q). Suppose that, almost surely, for every round r<Tr<T, client k, and local step 0≤e≤E0≤ e≤ E, Wkr,e∈Q(Wr)⊆,(Q).W_k^r,e _R_Q(W^r) _ γ, δ^(Q). Suppose also that finite deterministic envelopes MX,qM_X,q satisfy, almost surely for every round, ‖WXr,(q)‖2+Rq≤MX,q,X∈Δ,B,C,A,q∈1,…,Q, \|W_X^r,(q) \|_2+R_q≤ M_X,q, X∈\ ,B,C,A\, q∈\1,…,Q\, where WAr,(q):=Ar,(q)W_A^r,(q):=A^r,(q). The stability margins γq,δq>0 _q, _q>0, the layer-input bounds Lx,qL_x,q, the loss-regularity constants Lℓ,HℓL_ ,H_ , and the parameter envelopes MX,qM_X,q are deterministic and common across rounds. Let Assumption C.1 hold uniformly on all these neighborhoods, and for each layer, define ρq:=exp⁡(−γqδq)<1. _q:= (- _q _q)<1. Define: MKΔ,q M_K_ ,q :=MΔ,qLx,q+Dlog⁡(2), :=M_ ,qL_x,q+ D (2), MKB,q M_K_B,q :=MB,qLx,q2, :=M_B,qL_x,q^2, MKW,q M_K_W,q :=NDMKB,qMKΔ,q(MKΔ,q+DMA,qLx,q)(ρq1−ρq) :=NDM_K_B,qM_K_ ,q (M_K_ ,q+ DM_A,qL_x,q ) ( _q1- _q ) +Lx,q2NMKΔ,q+Lx,qDMKB,q, +L_x,q^2 NM_K_ ,q+L_x,q DM_K_B,q, HRq H_R_q :=MKB,qMKΔ,qND1−ρq, := M_K_B,qM_K_ ,q ND1- _q, JRq J_R_q :=MKW,q1−ρq, := M_K_W,q1- _q, MΛx,q M_ _x,q :=ρqMA,qMΔ,qHRq+MB,qLx,q(2MKΔ,q+Lx,qMΔ,q), := _qM_A,qM_ ,qH_R_q+M_B,qL_x,q (2M_K_ ,q+L_x,qM_ ,q ), Mκ,q M_κ,q :=MC,q(HRq+Lx,qMΛx,q1−ρq), :=M_C,q (H_R_q+ L_x,qM_ _x,q1- _q ), Θq _q :=(Lx,q2(HRq+MC,qJRq)2+Mκ,q2Θq−12)1/2,Θ0:=0, := (L_x,q^2 (H_R_q+M_C,qJ_R_q )^2+M_κ,q^2 _q-1^2 )^1/2, _0:=0, ℰq:= _q:= ρqHRq[(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))2+2(Lx,q+MΔ,qΘq−1) _qH_R_q [ (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) )^2+2 (L_x,q+M_ ,q _q-1 ) +MA,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)] +M_A,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 ) ] +2ρq(MKΔ,q+MA,q(Lx,q+MΔ,qΘq−1))(JRq+MΛx,qΘq−11−ρq) +2 _q (M_K_ ,q+M_A,q (L_x,q+M_ ,q _q-1 ) ) (J_R_q+ M_ _x,q _q-11- _q ) +2Θq−1MKΔ,qLx,q +2 _q-1M_K_ ,qL_x,q +2(Lx,q+MB,qΘq−1)[Lx,q(Lx,q+MΔ,qΘq−1)+MKΔ,qΘq−1] +2 (L_x,q+M_B,q _q-1 ) [L_x,q (L_x,q+M_ ,q _q-1 )+M_K_ ,q _q-1 ] +MB,qLx,q[Lx,q(14(Lx,q+MΔ,qΘq−1)2+2Θq−1)+2(Lx,q+MΔ,qΘq−1)Θq−1], +M_B,qL_x,q [L_x,q ( 14 (L_x,q+M_ ,q _q-1 )^2+2 _q-1 )+2 (L_x,q+M_ ,q _q-1 ) _q-1 ], Ωq _q :=MC,qLx,qℰq1−ρq+2(Lx,q+MC,qΘq−1)(JRq+MΛx,qΘq−11−ρq)+2HRqΘq−1+Mκ,qΩq−1, := M_C,qL_x,qE_q1- _q+2 (L_x,q+M_C,q _q-1 ) (J_R_q+ M_ _x,q _q-11- _q )+2H_R_q _q-1+M_κ,q _q-1, Ω0 _0 =0 =0 Let: LSSM,Q(Q) L_SSM,R_Q^(Q) :=HℓΘQ2+LℓΩQ, :=H_ _Q^2+L_ _Q, GSSM,Q(Q) G_SSM,R_Q^(Q) :=LℓΘQ=Lℓ(∑q=1QLx,q2(HRq+MC,qJRq)2∏q′=q+1QMκ,q′2)1/2, :=L_ _Q=L_ ( _q=1^QL_x,q^2 (H_R_q+M_C,qJ_R_q )^2 _q =q+1^QM_κ,q ^2 )^1/2, define the common deterministic smoothness and gradient bounds on all the considered neighborhoods. For the client distributions, define ¯:=∑k=1Kpkk,ℋTV:=∑k=1KpkTV⁡(k,¯)2, P:= _k=1^Kp_kP_k, _TV:= _k=1^Kp_kTV(P_k, P)^2, where TV⁡(P,Q):=sup|P⁡()−Q⁡()|.TV(P,Q):= _A|P(A)-Q(A)|. If the step size satisfies ηELSSM,Q(Q)≤12,12η2(LSSM,Q(Q))2E(E−1)≤1,η EL_SSM,R_Q^(Q)≤ 12, 12η^2 (L_SSM,R_Q^(Q) )^2E(E-1)≤ 1, then 1T∑r=0T−1‖∇f(Wr)‖F2≤ 1T _r=0^T-1E\|∇ f(W^r)\|_F^2≤ 4(f⁡(W0)−f∗)ηET+4ηLSSM,Q(Q)(GSSM,Q(Q))2Ξ∑k=1Kpk2 4(f(W^0)-f^*)η ET+ 4η L_SSM,R_Q^(Q) (G_SSM,R_Q^(Q) )^2 _k=1^Kp_k^2 +6η2(LSSM,Q(Q))2(E−1)(GSSM,Q(Q))2Ξ + 6η^2 (L_SSM,R_Q^(Q) )^2(E-1) (G_SSM,R_Q^(Q) )^2 +48η2(LSSM,Q(Q))2E(E−1)(GSSM,Q(Q))2ℋTV. +48η^2 (L_SSM,R_Q^(Q) )^2E(E-1) (G_SSM,R_Q^(Q) )^2H_TV. Proof of Corollary C.8: We apply the federated optimization argument from the proof of Theorem B.5 to the full parameter tuple W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ). After vectorization, the Euclidean norm of this tuple is exactly its block Frobenius norm. The preceding neighborhood estimates remain valid when the reference-dependent parameter envelopes are replaced by the common deterministic upper envelopes MX,qM_X,q, because their defining expressions are nondecreasing in those envelopes. Lemma C.4 therefore supplies the common smoothness constant L=LSSM,Q(Q)L=L_SSM,R_Q^(Q). Lemma C.6, applied conditionally on the history before each mini-batch, supplies the common variance bound σ2=(GSSM,Q(Q))2Ξ.σ^2= (G_SSM,R_Q^(Q) )^2 . Conditional independence across clients gives the corresponding weighted-noise variance bound σ2∑kpk2σ^2 _kp_k^2. Lemma C.7 supplies the common client-gradient disagreement bound Γ2=4(GSSM,Q(Q))2ℋTV. ^2=4 (G_SSM,R_Q^(Q) )^2H_TV. Each neighborhood Q(Wr)N_R_Q(W^r) is a product of convex balls. Since it contains all local iterates in round r, it also contains their weighted aggregate and the segments used in the smoothness argument. Substituting these constants into (164) yields the desired inequality in Corollary C.8 under the specified step size conditions. ∎ C.4 Corollary C.9: FedProx Multi-Layer SSM Convergence and Proof Corollary C.9 (FedProx convergence for multi-layer selective SSMs). Consider the distributed learning framework established in Section 3, with objective f⁡(W):=∑k=1Kpkℒk(W),pk≥0,∑k=1Kpk=1,f(W):= _k=1^Kp_kL_k(W), p_k≥ 0, _k=1^Kp_k=1, where the weights pkp_k are fixed and W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ) parameterizes the selective SSM of finite integer depth Q≥1Q≥ 1 defined in Appendix C.1. Each client objective is finite and differentiable on the neighborhoods under consideration and satisfies ∇Wℒk(W)=ξ∼k[∇WℒS,(Q)(W,ξ)], _WL_k(W)=E_ξ _k [ _WL^S,(Q)(W;ξ) ], where S is the length of the sampled sequence pair ξ, and all clients use the same sequence-loss map. Let W0W^0 be deterministic with f⁡(W0)<∞f(W^0)<∞, and suppose that f is bounded below by a finite constant f∗f^*. Consider full-participation FedProx for T communication rounds, with E local steps per round, mini-batch size Ξ , constant deterministic local step size η>0η>0, and constant deterministic proximal coefficient μ>0μ>0, where T,E,Ξ∈ℤ>0T,E, _>0: Wkr,0 W_k^r,0 =Wr, =W^r, Wkr,e+1 W_k^r,e+1 =Wkr,e−η(gk,Ξr,e+μ(Wkr,e−Wr)),e=0,…,E−1, =W_k^r,e-η (g_k, ^r,e+μ(W_k^r,e-W^r) ), e=0,…,E-1, Wr+1 W^r+1 =∑k=1KpkWkr,E,r=0,…,T−1. = _k=1^Kp_kW_k^r,E, r=0,…,T-1. The mini-batch gradient is gk,Ξr,e:=1Ξ∑i=1Ξ∇WℒSi,(Q)(Wkr,e,ξk,ir,e),g_k, ^r,e:= 1 _i=1 _WL^S_i,(Q) (W_k^r,e; _k,i^r,e ), where SiS_i denotes the length of ξk,ir,e _k,i^r,e. Conditional on the complete training history before each local step, each client draws Ξ fresh independent samples from kP_k, and the mini-batches drawn by different clients are conditionally independent. For fixed deterministic radii Rq>0R_q>0, define Q(Wr):=∏q=1QRq(q)(Wr,(q)),Q:=(R1,…,RQ).N_R_Q(W^r):= _q=1^QN_R_q^(q)(W^r,(q)), _Q:=(R_1,…,R_Q). Suppose that, almost surely, for every round r<Tr<T, client k, and local step 0≤e≤E0≤ e≤ E, Wkr,e∈Q(Wr)⊆,(Q).W_k^r,e _R_Q(W^r) _ γ, δ^(Q). Suppose also that finite deterministic envelopes MX,qM_X,q satisfy, almost surely for every round, ‖WXr,(q)‖2+Rq≤MX,q,X∈Δ,B,C,A,q=1,…,Q, \|W_X^r,(q) \|_2+R_q≤ M_X,q, X∈\ ,B,C,A\, q=1,…,Q, where WAr,(q):=Ar,(q)W_A^r,(q):=A^r,(q). Let Assumption C.1 hold uniformly on all these neighborhoods, and let Assumption 4.2 hold uniformly at the final-layer predictions. The stability margins γq,δq>0 _q, _q>0, layer-input bounds Lx,qL_x,q, loss-regularity constants Lℓ,HℓL_ ,H_ , and parameter envelopes MX,qM_X,q are deterministic and common across rounds, admissible sequence pairs, finite sequence lengths, and token positions. For each layer, define ρq:=exp⁡(−γqδq)<1. _q:= (- _q _q)<1. Let GSSM,Q(Q)G_SSM,R_Q^(Q) be the common gradient bound defined in Corollary (C.8), constructed using the original envelopes MX,qM_X,q. Assume 0<ημ≤10<ημ≤ 1, and define Rμ,q R_μ,q :=minRq,GSSM,Q(Q)μ[1−(1−ημ)E],q=1,…,Q, := \R_q, G_SSM,R_Q^(Q)μ [1-(1-ημ)^E ] \, q=1,…,Q, Q,μ _Q,μ :=(Rμ,1,…,Rμ,Q), :=(R_μ,1,…,R_μ,Q), MX,qμ M_X,q^μ :=MX,q−Rq+Rμ,q,X∈Δ,B,C,A. :=M_X,q-R_q+R_μ,q, X∈\ ,B,C,A\. Using these reduced parameter envelopes, define MKΔ,qμ M_K_ ,q^μ :=MΔ,qμLx,q+Dlog⁡2, :=M_ ,q^μL_x,q+ D 2, MKB,qμ M_K_B,q^μ :=MB,qμLx,q2, :=M_B,q^μL_x,q^2, MKW,qμ M_K_W,q^μ :=NDMKB,qμMKΔ,qμ(MKΔ,qμ+DMA,qμLx,q)ρq1−ρq :=NDM_K_B,q^μM_K_ ,q^μ (M_K_ ,q^μ+ DM_A,q^μL_x,q ) _q1- _q +Lx,q2NMKΔ,qμ+Lx,qDMKB,qμ, +L_x,q^2 NM_K_ ,q^μ+L_x,q DM_K_B,q^μ, HRμ,q H_R_μ,q :=MKB,qμMKΔ,qμND1−ρq, := M_K_B,q^μM_K_ ,q^μ ND1- _q, JRμ,q J_R_μ,q :=MKW,qμ1−ρq, := M_K_W,q^μ1- _q, MΛx,qμ M_ _x,q^μ :=ρqMA,qμMΔ,qμHRμ,q+MB,qμLx,q(2MKΔ,qμ+Lx,qMΔ,qμ), := _qM_A,q^μM_ ,q^μH_R_μ,q+M_B,q^μL_x,q (2M_K_ ,q^μ+L_x,qM_ ,q^μ ), Mκ,qμ M_κ,q^μ :=MC,qμ(HRμ,q+Lx,qMΛx,qμ1−ρq). :=M_C,q^μ (H_R_μ,q+ L_x,qM_ _x,q^μ1- _q ). Initialize Θ0μ=Ω0μ=0 _0^μ= _0^μ=0, and recursively define Θqμ _q^μ :=[Lx,q2(HRμ,q+MC,qμJRμ,q)2+(Mκ,qμ)2(Θq−1μ)2]1/2, := [L_x,q^2 (H_R_μ,q+M_C,q^μJ_R_μ,q )^2+(M_κ,q^μ)^2( _q-1^μ)^2 ]^1/2, ℰqμ _q^μ :=ρqHRμ,q[(MKΔ,qμ+MA,qμ(Lx,q+MΔ,qμΘq−1μ))2 := _qH_R_μ,q [ (M_K_ ,q^μ+M_A,q^μ (L_x,q+M_ ,q^μ _q-1^μ ) )^2 +2(Lx,q+MΔ,qμΘq−1μ)+MA,qμ(14(Lx,q+MΔ,qμΘq−1μ)2+2Θq−1μ)] +2 (L_x,q+M_ ,q^μ _q-1^μ )+M_A,q^μ ( 14 (L_x,q+M_ ,q^μ _q-1^μ )^2+2 _q-1^μ ) ] +2ρq(MKΔ,qμ+MA,qμ(Lx,q+MΔ,qμΘq−1μ))(JRμ,q+MΛx,qμΘq−1μ1−ρq) +2 _q (M_K_ ,q^μ+M_A,q^μ (L_x,q+M_ ,q^μ _q-1^μ ) ) (J_R_μ,q+ M_ _x,q^μ _q-1^μ1- _q ) +2Θq−1μMKΔ,qμLx,q +2 _q-1^μM_K_ ,q^μL_x,q +2(Lx,q+MB,qμΘq−1μ)[Lx,q(Lx,q+MΔ,qμΘq−1μ)+MKΔ,qμΘq−1μ] +2 (L_x,q+M_B,q^μ _q-1^μ ) [L_x,q (L_x,q+M_ ,q^μ _q-1^μ )+M_K_ ,q^μ _q-1^μ ] +MB,qμLx,q[Lx,q(14(Lx,q+MΔ,qμΘq−1μ)2+2Θq−1μ) +M_B,q^μL_x,q [L_x,q ( 14 (L_x,q+M_ ,q^μ _q-1^μ )^2+2 _q-1^μ ) +2(Lx,q+MΔ,qμΘq−1μ)Θq−1μ], +2 (L_x,q+M_ ,q^μ _q-1^μ ) _q-1^μ ], Ωqμ _q^μ :=MC,qμLx,qℰqμ1−ρq := M_C,q^μL_x,qE_q^μ1- _q +2(Lx,q+MC,qμΘq−1μ)(JRμ,q+MΛx,qμΘq−1μ1−ρq) +2 (L_x,q+M_C,q^μ _q-1^μ ) (J_R_μ,q+ M_ _x,q^μ _q-1^μ1- _q ) +2HRμ,qΘq−1μ+Mκ,qμΩq−1μ. +2H_R_μ,q _q-1^μ+M_κ,q^μ _q-1^μ. Let LSSM,Q,μ(Q) L_SSM,R_Q,μ^(Q) :=Hℓ(ΘQμ)2+LℓΩQμ, :=H_ ( _Q^μ)^2+L_ _Q^μ, Lprox,Q,μ(Q) L_prox,R_Q,μ^(Q) :=LSSM,Q,μ(Q)+μ, :=L_SSM,R_Q,μ^(Q)+μ, GSSM,Q,μ(Q) G_SSM,R_Q,μ^(Q) :=LℓΘQμ :=L_ _Q^μ =Lℓ[∑q=1QLx,q2(HRμ,q+MC,qμJRμ,q)2∏q′=q+1Q(Mκ,q′μ)2]1/2, =L_ [ _q=1^QL_x,q^2 (H_R_μ,q+M_C,q^μJ_R_μ,q )^2 _q =q+1^Q(M_κ,q ^μ)^2 ]^1/2, where empty products equal one. These are common deterministic smoothness and gradient bounds on the reduced neighborhoods Q,μ(Wr):=∏q=1QRμ,q(q)(Wr,(q)).N_R_Q,μ(W^r):= _q=1^QN_R_μ,q^(q)(W^r,(q)). For the client distributions, define ¯:=∑k=1Kpkk,ℋTV:=∑k=1KpkTV⁡(k,¯)2, P:= _k=1^Kp_kP_k, _TV:= _k=1^Kp_kTV(P_k, P)^2, where TV⁡(P,Q):=sup|P⁡()−Q⁡()|TV(P,Q):= _A|P(A)-Q(A)|, with the supremum taken over measurable sets A. If the step size satisfies ηELSSM,Q,μ(Q)≤12,12η2(Lprox,Q,μ(Q))2E(E−1)≤1,η EL_SSM,R_Q,μ^(Q)≤ 12, 12η^2 (L_prox,R_Q,μ^(Q) )^2E(E-1)≤ 1, then 1T∑r=0T−1‖∇f(Wr)‖F2≤ 1T _r=0^T-1E\|∇ f(W^r)\|_F^2≤ 3(f⁡(W0)−f∗)ηET 3(f(W^0)-f^*)η ET +3ηLSSM,Q,μ(Q)(GSSM,Q,μ(Q))2Ξ∑k=1Kpk2 + 3η L_SSM,R_Q,μ^(Q) (G_SSM,R_Q,μ^(Q) )^2 _k=1^Kp_k^2 +3η2(Lprox,Q,μ(Q))2(E−1)(GSSM,Q,μ(Q))2Ξ + 3η^2 (L_prox,R_Q,μ^(Q) )^2(E-1) (G_SSM,R_Q,μ^(Q) )^2 +24η2(Lprox,Q,μ(Q))2E(E−1)(GSSM,Q,μ(Q))2ℋTV. +24η^2 (L_prox,R_Q,μ^(Q) )^2E(E-1) (G_SSM,R_Q,μ^(Q) )^2H_TV. Moreover, almost surely, for every round r<Tr<T, client k, layer q, and local step 0≤e≤E0≤ e≤ E, ‖Wkr,e,(q)−Wr,(q)‖F≤min⁡Rq,GSSM,Q(Q)μ[1−(1−ημ)e]≤Rμ,q. \|W_k^r,e,(q)-W^r,(q) \|_F≤ \R_q, G_SSM,R_Q^(Q)μ [1-(1-ημ)^e ] \≤ R_μ,q. Consequently, all local iterates and the aggregated model Wr+1W^r+1 lie in Q,μ(Wr)⊆Q(Wr)⊆,(Q).N_R_Q,μ(W^r) _R_Q(W^r) _ γ, δ^(Q). Proof of Corollary C.9: We apply the federated optimization argument from the proof of Theorem B.6 to the full parameter tuple W=(W(1),…,W(Q))W= (W^(1),…,W^(Q) ). After vectorization, the Euclidean norm of this tuple is exactly its block Frobenius norm. First, the pointwise multi-layer gradient bound implies ‖gk,Ξr,e‖F≤GSSM,Q(Q) \|g_k, ^r,e \|_F≤ G_SSM,R_Q^(Q), which also bounds each layer block of the mini-batch gradient. Applying the displacement recursion from the single-layer FedProx proof separately to each layer, and using 0<ημ≤10<ημ≤ 1, gives ‖Wkr,e,(q)−Wr,(q)‖F≤GSSM,Q(Q)μ[1−(1−ημ)e]. \|W_k^r,e,(q)-W^r,(q) \|_F≤ G_SSM,R_Q^(Q)μ [1-(1-ημ)^e ]. (403) Combining this with the assumed containment in the original neighborhood yields the stated bound by Rμ,qR_μ,q. Consequently, all local iterates lie in Q,μ(Wr)N_R_Q,μ(W^r). This neighborhood is a product of convex balls, so it also contains their weighted aggregate and the segments used in the smoothness argument. For every W∈Q,μ(Wr)W _R_Q,μ(W^r), ‖WX(q)‖2≤‖WXr,(q)‖2+Rμ,q≤MX,q−Rq+Rμ,q=MX,qμ. \|W_X^(q) \|_2≤ \|W_X^r,(q) \|_2+R_μ,q≤ M_X,q-R_q+R_μ,q=M_X,q^μ. The preceding neighborhood estimates therefore apply with the common deterministic envelopes MX,qμM_X,q^μ, because their defining expressions are nondecreasing in those envelopes. Lemma C.4 supplies the common smoothness constant LSSM,Q,μ(Q)L_SSM,R_Q,μ^(Q) for each client objective and for f. Adding the proximal penalty makes each local proximal objective Lprox,Q,μ(Q)L_prox,R_Q,μ^(Q)-smooth, where Lprox,Q,μ(Q)=LSSM,Q,μ(Q)+μ.L_prox,R_Q,μ^(Q)=L_SSM,R_Q,μ^(Q)+μ. Lemma C.6, applied conditionally on the complete history before each mini-batch, supplies σ2=(GSSM,Q,μ(Q))2Ξ.σ^2= (G_SSM,R_Q,μ^(Q) )^2 . The proximal term is deterministic under this conditioning and does not alter the gradient-noise variance. Conditional independence across clients gives the corresponding weighted-noise variance bound σ2∑kpk2σ^2 _kp_k^2. Lemma C.7 supplies the common client-gradient disagreement bound Γ2=4(GSSM,Q,μ(Q))2ℋTV. ^2=4 (G_SSM,R_Q,μ^(Q) )^2H_TV. The common proximal term cancels when comparing client gradients with their weighted average at the same parameter tuple. The remaining local-drift and descent calculations in the proof of Theorem B.6 depend only on these smoothness, variance, and disagreement bounds, the conditional sampling assumptions, and convexity of the neighborhoods. Applying those calculations with the multi-layer constants above under the state learning-rate conditions and substituting the defined values of σ2σ^2 and Γ2 ^2 yields the desired inequality in Corollary C.9, completing the proof. ∎ Appendix D Supplementary Corollaries In this section, we provide supplementary corollaries to the theorems presented in the preceding sections and in the main manuscript. These corollaries simplify the preceding results and make explicit their dependence on SSM-specific quantities, including stability, parameter norms, network depth, and related architectural properties. D.1 Single-layer convergence rate corollaries In this section, we provide convergence rate guarantees for FedAvg [McMahan et al., 2017] and FedProx [Li et al., 2020] learning frameworks. Every convergence corollary retains the assumptions of its referenced result. Furthermore, we establish and retain the following conventions: F0 F_0 :=f⁡(W0)−f∗, :=f(W^0)-f^*, (404) P2 P_2 :=∑k=1Kpk2, := _k=1^Kp_k^2, (405) ℋTV _TV :=∑kpkTV⁡(k,¯)2 := _kp_kTV(P_k, P)^2 (406) All O expressions denote upper bounds on nonnegative quantities. Note that Corollary 4.8 is stated in Section 4, with this section containing its proof. D.1.1 Corollary 4.8: Single-layer FedAvg convergence rate Proof Proof of Corollary 4.8: By Lemmas 4.7 and B.2: GSSM,R=⁡(MR(1−ρ)2),LSSM,R=⁡(MR2(1−ρ)4)G_SSM,R=O ( M_R(1-ρ)^2 ), L_SSM,R=O ( M_R^2(1-ρ)^4 ) (407) Substituting these bounds into Theorem B.5, collecting the two local-drift terms, and using the fact that P2≤1P_2≤ 1, yields the inequality in Corollary 4.8. ∎ D.1.2 Corollary D.1: Single-layer FedProx convergence rate and Proof Corollary D.1 (Single-layer FedProx convergence rate). Under the assumptions and step-size conditions of Theorem B.6, given its reduced envelopes MXμ=MX−R+RμM_X^μ=M_X-R+R_μ, define: MRμ:=(1+MBμ)(1+MCμ)(1+MΔμ)(1+MΔμ+MAμ).M_R_μ:=(1+M_B^μ)(1+M_C^μ)(1+M_ ^μ)(1+M_ ^μ+M_A^μ). Then 1T∑r=0T−1‖∇f(Wr)‖F2=⁡(F0ηET+ηMRμ4Ξ(1−ρ)8+η2MRμ2(1−ρ)4(MRμ2(1−ρ)4+μ)2(E−1)(1Ξ+EℋTV)). 1T _r=0^T-1E \|∇ f(W^r) \|_F^2=O ( F_0η ET+ η M_R_μ^4 (1-ρ)^8+η^2 M_R_μ^2(1-ρ)^4 ( M_R_μ^2(1-ρ)^4+μ )^2(E-1) ( 1 +EH_TV ) ). (408) The implicit constant has the same dependencies as in Corollary 4.8 and is independent of μ. Proof of Corollary D.1: The single-layer scaling argument, evaluated with MXμM_X^μ, gives GSSM,Rμ=⁡(MRμ(1−ρ)2)G_SSM,R_μ=O ( M_R_μ(1-ρ)^2 ) and LSSM,Rμ=⁡(MRμ2(1−ρ)4)L_SSM,R_μ=O ( M_R_μ^2(1-ρ)^4 ). Consequently Lprox,Rμ=⁡(MRμ2(1−ρ)4+μ)L_prox,R_μ=O ( M_R_μ^2(1-ρ)^4+μ ). Substituting these results into Theorem B.6 and using the fact that P2≤1P_2≤ 1 proves the bound. ∎ Note that although MRμ(1−ρ)2≤MR(1−ρ)2 M_R_μ(1-ρ)^2≤ M_R(1-ρ)^2 for the same original envelopes, the bound in Corollary D.1 also contains additional μ terms. As such, these bounds expose a trade-off and do not imply that every positive proximal coefficient improves the convergence guarantee. D.2 Multi-Layer SSM Gradient Scaling Bounds In this section we provide simplified bounds on the multi-layer gradient, where all per-layer quantities are scaled to a single, common value. For the remainder of this section, we will utilize the following notation: for a fixed model in Theorem C.3, let Lx L_x :=max1≤q≤QLx,q,ρ:=max1≤q≤Qρq<1, := _1≤ q≤ QL_x,q, ρ:= _1≤ q≤ Q _q<1, (409) MX M_X :=max1≤q≤Q‖WX(q)‖2,X∈Δ,B,C,A, := _1≤ q≤ Q \|W_X^(q) \|_2, X∈\ ,B,C,A\, where WA(q):=A(q)W_A^(q):=A^(q). For neighborhood-uniform statements, take MX:=maxq⁡MX,qM_X:= _qM_X,q using the deterministic envelopes in the multi-layer convergence results. The bounds on LxL_x and ρ are then uniform on those neighborhoods as well. Note that common upper bounds suffice - the layers do not need to have identical weights or identical stability factors. Using these common quantities, define MKΔ M_K_ :=MΔLx+Dlog⁡2, :=M_ L_x+ D 2, (410) MKB M_K_B :=MBLx2, :=M_BL_x^2, (411) MKW M_K_W :=NDMKBMKΔ(MKΔ+DMALx)ρ1−ρ :=NDM_K_BM_K_ (M_K_ + DM_AL_x) ρ1-ρ +Lx2NMKΔ+LxDMKB, +L_x^2 NM_K_ +L_x DM_K_B, (412) H H :=MKBMKΔND1−ρ, := M_K_BM_K_ ND1-ρ, (413) J J :=MKW1−ρ, := M_K_W1-ρ, (414) Λ :=ρMAMΔH+MBLx(2MKΔ+LxMΔ), :=ρ M_AM_ H+M_BL_x(2M_K_ +L_xM_ ), (415) κ κ :=MC(H+LxΛ1−ρ). :=M_C (H+ L_x 1-ρ ). (416) D.2.1 Corollary D.2: Multi-layer gradient under common layer bounds and Proof Corollary D.2 (Multi-layer gradient under common layer bounds). Under the assumptions of Theorem C.3, the quantities in (410)-(416) satisfy, for every admissible sequence pair ξ of finite length S, ‖∇WℒS,(Q)(W,ξ)‖F≤LℓLx(H+MCJ)∑q=0Q−1κ2q. \| _WL^S,(Q)(W;ξ) \|_F≤ L_ L_x(H+M_CJ) _q=0^Q-1κ^2q. Furthermore, the following inequalities hold for different values of κ: ‖∇WℒS,(Q)(W,ξ)‖F≤LℓLx(H+MCJ)11−κ2,0≤κ<1,Q,κ=1,(κ2Q−1)/(κ2−1),κ>1. \| _WL^S,(Q)(W;ξ) \|_F≤ L_ L_x(H+M_CJ) cases 1 1-κ^2,&0≤κ<1,\\ Q,&κ=1,\\ (κ^2Q-1)/(κ^2-1),&κ>1. cases Proof of Corollary D.2: All defining expressions 410- 416 are nondecreasing in the nonnegative parameter, input, and stability upper bounds. As a result, each own-layer factor in Theorem C.3 is at most Lx(H+MCJ)L_x(H+M_CJ), and each downstream sensitivity is at most κ. Substitution into the inequality from Theorem C.3 gives the right-hand side bounded by LℓLx(H+MCJ)(∑q=1Qκ2(Q−q))1/2L_ L_x(H+M_CJ) ( _q=1^Qκ^2(Q-q) )^1/2. Evaluating the geometric sum proves both claimed inequalities. ∎ For common envelopes independent of Q, the upper bound is uniform in depth when κ<1κ<1, grows at most as Q Q when κ=1κ=1, and can grow exponentially when κ>1κ>1. Temporal stability ρ<1ρ<1 alone does not imply κ<1κ<1, and an exponentially growing upper bound does not establish that actual gradients grow exponentially. D.2.2 Corollary D.3: Multi-layer gradient rate under common layer bounds and Proof Corollary D.3 (Multi-layer gradient rate under common layer bounds). Under the assumptions of Theorem C.3, define q:=(1+‖WB(q)‖2)(1+‖WC(q)‖2)(1+‖WΔ(q)‖2)(1+‖WΔ(q)‖2+‖A(q)‖2). M_q:= (1+\|W_B^(q)\|_2 ) (1+\|W_C^(q)\|_2 ) (1+\|W_ ^(q)\|_2 ) (1+\|W_ ^(q)\|_2+\|A^(q)\|_2 ). Then, given the κq _q defined in that theorem, it holds that: ‖∇WℒS,(Q)(W,ξ)‖F=⁡[(∑q=1Qq2(1−ρq)4∏j=q+1Qκj2)1/2]. \| _WL^S,(Q)(W;ξ) \|_F=O [ ( _q=1^Q M_q^2(1- _q)^4 _j=q+1^Q _j^2 )^1/2 ]. Furthermore, with the common envelopes defined above, let :=(1+MB)(1+MC)(1+MΔ)(1+MΔ+MA) M:=(1+M_B)(1+M_C)(1+M_ )(1+M_ +M_A). Then the following holds: ‖∇WℒS,(Q)(W,ξ)‖F \| _WL^S,(Q)(W;ξ) \|_F =⁡((1−ρ)2∑q=0Q−1κ2q) =O ( M(1-ρ)^2 _q=0^Q-1κ^2q ) =⁡((1−ρ)211−κ2,0≤κ<1,Q,κ=1,(κ2Q−1)/(κ2−1),κ>1.). =O ( M(1-ρ)^2 cases 1 1-κ^2,&0≤κ<1,\\ Q,&κ=1,\\ (κ^2Q-1)/(κ^2-1),&κ>1. cases ). The implicit constants depend only on N,D,Lx,LℓN,D,L_x,L_ , where Lx≥maxq⁡Lx,qL_x≥ _qL_x,q, and are independent of Q, sequence length, parameter norms, and stability factors. Proof of Corollary D.3: The single-layer estimates give Lx,q(Hq+‖WC(q)‖2Jq)≤Cq/(1−ρq)2L_x,q (H_q+ \|W_C^(q) \|_2J_q )≤ C M_q/(1- _q)^2 for some common constant C. Substituting this into Theorem C.3 we obtain the first equality of in Corollary D.3. Additionally, substituting Lx(H+MCJ)=⁡(/(1−ρ)2)L_x(H+M_CJ)=O ( M/(1-ρ)^2 ) into Corollary D.2 proves the second equality. ∎ D.3 Multi-layer convergence under common layer bounds In the following corollaries and their proof, we utilize the neighborhood-uniform common envelopes defined in (409) and 410- 416, with the same Lℓ,HℓL_ ,H_ throughout all considered neighborhoods. The following reduction replaces the original layer-dependent sensitivity recursions by three finite geometric sums. We further define the common auxiliary constants d∗ d_* :=Lx+MΔ, :=L_x+M_ , (417) b∗ b_* :=14d∗2+2, := 14d_*^2+2, (418) c∗ c_* :=MKΔ+MAd∗, :=M_K_ +M_Ad_*, (419) j∗ j_* :=J+Λ1−ρ, :=J+ 1-ρ, (420) ℰ∗ _* :=ρH(c∗2+2d∗+MAb∗)+2ρc∗j∗+2MKΔLx :=ρ H(c_*^2+2d_*+M_Ab_*)+2ρ c_*j_*+2M_K_ L_x +2(Lx+MB)(Lxd∗+MKΔ)+MBLx(Lxb∗+2d∗), +2(L_x+M_B)(L_xd_*+M_K_ )+M_BL_x(L_xb_*+2d_*), (421) B∗ B_* :=MCLxℰ∗1−ρ+2(Lx+MC)j∗+2H. := M_CL_xE_*1-ρ+2(L_x+M_C)j_*+2H. (422) For Q≥1Q≥ 1, define G^Q G_Q :=LℓLx(H+MCJ)∑q=0Q−1κ2q, :=L_ L_x(H+M_CJ) _q=0^Q-1κ^2q, (423) L^Q L_Q :=HℓLx2(H+MCJ)2∑q=0Q−1κ2q :=H_ L_x^2(H+M_CJ)^2 _q=0^Q-1κ^2q +LℓB∗[∑q=0Q−1κq+Lx2(H+MCJ)2∑q=1QκQ−q(∑j=0q−1κ2j)]. +L_ B_* [ _q=0^Q-1κ^q+L_x^2(H+M_CJ)^2 _q=1^Qκ^Q-q ( _j=0^q-1κ^2j ) ]. (424) In particular, for κ≠1κ≠ 1, ∑q=0Q−1κq=1−κQ1−κ,∑q=1QκQ−q(∑j=0q−1κ2j)=(1−κQ)(1−κQ+1)(1−κ)(1−κ2),κ≠1. _q=0^Q-1κ^q= 1-κ^Q1-κ, _q=1^Qκ^Q-q ( _j=0^q-1κ^2j )= (1-κ^Q)(1-κ^Q+1)(1-κ)(1-κ^2), κ≠ 1. (425) For κ=1κ=1, the two sums reduce to Q and Q⁡(Q+1)/2Q(Q+1)/2, respectively. D.3.1 Corollary D.4: Common-layer smoothness and depth dependence and Proof Corollary D.4 (Common-layer smoothness and depth dependence). Under the neighborhood assumptions of Lemma C.4, the constants GSSM,Q(Q)G_SSM,R_Q^(Q) and LSSM,Q(Q)L_SSM,R_Q^(Q) defined in Corollary C.8 satisfy: GSSM,Q(Q)≤G^Q,LSSM,Q(Q)≤L^Q.G_SSM,R_Q^(Q)≤ G_Q, L_SSM,R_Q^(Q)≤ L_Q. If all common envelopes and regularity constants are fixed independently of Q, their depth dependence is (G^Q,L^Q)=(⁡(1),⁡(1)),0≤κ<1,(⁡(Q),⁡(Q2)),κ=1,(⁡(κQ),⁡(κ2Q)),κ>1. ( G_Q, L_Q )= cases (O(1),O(1) ),&0≤κ<1,\\ (O( Q),O(Q^2) ),&κ=1,\\ (O(κ^Q),O(κ^2Q) ),&κ>1. cases Here only Q varies: the implicit constants may depend on the fixed common envelopes and on κ. Proof of Corollary D.4: After replacing the layerwise quantities by their common upper bounds, the first-sensitivity recursion yields, for every q≥1q≥ 1, Θq2≤Lx2(H+MCJ)2∑j=0q−1κ2j. _q^2≤ L_x^2(H+M_CJ)^2 _j=0^q-1κ^2j. (426) Hence, GSSM,Q(Q)=LℓΘQ≤LℓLx(H+MCJ)∑j=0Q−1κ2j=G^Q.G_SSM,R_Q^(Q)=L_ _Q≤ L_ L_x(H+M_CJ) _j=0^Q-1κ^2j= G_Q. (427) For the second-sensitivity recursion, the terms independent of Ωq−1 _q-1 are bounded by B∗(1+Θq−12),B_* (1+ _q-1^2 ), (428) This follows because, after replacing the layerwise quantities by their common upper bounds, the terms independent of Ωq−1 _q-1 form a polynomial of degree at most two in Θq−1 _q-1 with nonnegative coefficients. The sum of these coefficients is B∗B_*, obtained by setting Θq−1=1 _q-1=1. Since each of 11, Θq−1 _q-1, and Θq−12 _q-1^2 is at most 1+Θq−121+ _q-1^2, the displayed bound follows. Therefore, unrolling the recursion gives ΩQ _Q ≤B∗∑q=1QκQ−q(1+Θq−12) ≤ B_* _q=1^Qκ^Q-q (1+ _q-1^2 ) ≤B∗[∑q=0Q−1κq+Lx2(H+MCJ)2∑q=1QκQ−q(∑j=0q−1κ2j)], ≤ B_* [ _q=0^Q-1κ^q+L_x^2(H+M_CJ)^2 _q=1^Qκ^Q-q ( _j=0^q-1κ^2j ) ], (429) where the first-sensitivity bound was used in the second inequality. Substituting the bounds on ΘQ _Q and ΩQ _Q into LSSM,Q(Q)=HℓΘQ2+LℓΩQ,L_SSM,R_Q^(Q)=H_ _Q^2+L_ _Q, (430) we get: LSSM,Q(Q)≤L^Q.L_SSM,R_Q^(Q)≤ L_Q. (431) Finally, evaluating the geometric sums gives (G^Q,L^Q)=(⁡(1),⁡(1)),0≤κ<1,(⁡(Q),⁡(Q2)),κ=1,(⁡(κQ),⁡(κ2Q)),κ>1, ( G_Q, L_Q )= cases (O(1),O(1) ),&0≤κ<1,\\ (O( Q),O(Q^2) ),&κ=1,\\ (O(κ^Q),O(κ^2Q) ),&κ>1, cases (432) which proves the stated depth dependence. ∎ D.3.2 Corollary D.5: Multi-layer FedAvg under common layer bounds and Proof Corollary D.5 (Multi-layer FedAvg under common layer bounds). Under the assumptions of Corollary C.8, let G^Q G_Q and L^Q L_Q be the common upper bounds in (423) and (424), evaluated with the neighborhood-uniform common envelopes introduced above. If ηEL^Q≤12,12η2L^Q2E(E−1)≤1,η E L_Q≤ 12, 12η^2 L_Q^2E(E-1)≤ 1, then 1T∑r=0T−1‖∇f(Wr)‖F2=⁡(CLOSE 1T _r=0^T-1E \|∇ f(W^r) \|_F^2=O ( F0ηET+ηL^QG^Q2P2Ξ F_0η ET+ η L_Q G_Q^2P_2 OPEN+η2L^Q2G^Q2(E−1)[1Ξ+EℋTV]). +η^2 L_Q^2 G_Q^2(E-1) [ 1 +EH_TV ] ). The implicit constant is numerical. For common envelopes and loss-regularity constants fixed independently of Q, the depth dependence of G^Q G_Q and L^Q L_Q is explicit in (423) and (424) and Corollary D.4. Proof of Corollary D.5: Corollary D.4 gives GSSM,Q(Q)≤G^Q,LSSM,Q(Q)≤L^Q.G_SSM,R_Q^(Q)≤ G_Q, L_SSM,R_Q^(Q)≤ L_Q. Consequently, the displayed step-size conditions imply those of Corollary C.8. Substituting these upper bounds into that corollary and grouping the two local-drift terms proves the stated claim. ∎ D.3.3 Corollary D.6: Multi-layer FedProx under common layer bounds and Proof Corollary D.6 (Multi-layer FedProx under common layer bounds). Under the assumptions of Corollary C.9, retain its original radii RqR_q, effective radii Rμ,qR_μ,q, and reduced envelopes MX,qμ=MX,q−Rq+Rμ,qM_X,q^μ=M_X,q-R_q+R_μ,q. For this corollary and its proof, evaluate the existing common bounds G^Q G_Q and L^Q L_Q in (423) and (424) using Lx L_x :=max1≤q≤QLx,q,ρ:=max1≤q≤Qρq, := _1≤ q≤ QL_x,q, ρ:= _1≤ q≤ Q _q, MX M_X :=max1≤q≤QMX,qμ,X∈Δ,B,C,A. := _1≤ q≤ QM_X,q^μ, X∈\ ,B,C,A\. All auxiliary quantities in (410)-(416) and (417)-(422), including κ and B∗B_*, are evaluated with these reduced common envelopes. The loss-regularity constants Lℓ,HℓL_ ,H_ are unchanged. If 0<ημ≤10<ημ≤ 1 and ηEL^Q≤12,12η2(L^Q+μ)2E(E−1)≤1,η E L_Q≤ 12, 12η^2 ( L_Q+μ )^2E(E-1)≤ 1, then 1T∑r=0T−1‖∇f(Wr)‖F2=⁡(CLOSE 1T _r=0^T-1E \|∇ f(W^r) \|_F^2=O ( F0ηET+ηL^QG^Q2P2Ξ F_0η ET+ η L_Q G_Q^2P_2 OPEN+η2(L^Q+μ)2G^Q2(E−1)[1Ξ+EℋTV]). +η^2( L_Q+μ)^2 G_Q^2(E-1) [ 1 +EH_TV ] ). The implicit constant is numerical and independent of μ. Proof of Corollary D.6: The displacement argument in the proof of Corollary C.9 places all local iterates in Q,μ(Wr)N_R_Q,μ(W^r). This product of convex balls also contains their weighted aggregate and the segments used in the smoothness argument. Since MX,qμ≤MXM_X,q^μ≤ M_X for every layer, the common-envelope construction of Corollary D.4 gives GSSM,Q,μ(Q)≤G^Q,LSSM,Q,μ(Q)≤L^Q.G_SSM,R_Q,μ^(Q)≤ G_Q, L_SSM,R_Q,μ^(Q)≤ L_Q. Therefore, Lprox,Q,μ(Q)=LSSM,Q,μ(Q)+μ≤L^Q+μ.L_prox,R_Q,μ^(Q)=L_SSM,R_Q,μ^(Q)+μ≤ L_Q+μ. The displayed step-size conditions consequently imply those of Corollary C.9. Substituting these upper bounds into that corollary and grouping the two local-drift terms proves the stated claim. ∎ D.4 SSM Convergence under Distributed Learning In this section we provide and prove Corollary D.7: SSM Convergence under Distributed Learning and the corresponding rate corollary - a setting where every client runs only one step of gradient descent between model weights exchange. Corollary D.7 (SSM Convergence under Distributed Learning). Under the assumptions of Theorem B.5, let E=1E=1. If ηLSSM,R≤1/2η L_SSM,R≤ 1/2, then 1T∑r=0T−1‖∇f(Wr)‖F2≤4F0ηT+4ηLSSM,RGSSM,R2P2Ξ. 1T _r=0^T-1E \|∇ f(W^r) \|_F^2≤ 4F_0η T+ 4η L_SSM,RG_SSM,R^2P_2 . Similarly, under the assumptions of Corollary C.8 with E=1E=1, if ηLSSM,Q(Q)≤1/2η L_SSM,R_Q^(Q)≤ 1/2, then 1T∑r=0T−1‖∇f(Wr)‖F2≤4F0ηT+4ηLSSM,Q(Q)(GSSM,Q(Q))2P2Ξ. 1T _r=0^T-1E \|∇ f(W^r) \|_F^2≤ 4F_0η T+ 4η L_SSM,R_Q^(Q) (G_SSM,R_Q^(Q) )^2P_2 . Both bounds also apply to FedProx with one local step, under the same FedAvg assumptions and step-size conditions, for every μ>0μ>0. No restriction on ημημ is needed for this one-step conclusion. Proof of Corollary D.7: Setting E=1E=1 in Theorem B.5 or Corollary C.8 eliminates both terms containing E−1E-1 and makes the second step-size condition automatic. For FedProx, initialization at the broadcast model gives Wkr,1=Wr−η⁡(gk,Ξr,0+μ⁡(Wr−Wr))=Wr−ηgk,Ξr,0.W_k^r,1=W^r-η (g_k, ^r,0+μ(W^r-W^r) )=W^r-η g_k, ^r,0. (433) As such, for every μ>0μ>0, its one-step update and weighted aggregate coincide with those of FedAvg. Applying the corresponding FedAvg result proves the claim. ∎ Corollary D.8 (SSM Convergence Rate under Distributed Learning). Under the assumptions of Corollary 4.8, let E=1E=1. Under ηLSSM,R≤1/2η L_SSM,R≤ 1/2, let MB,MC,MΔ,M_B,M_C,M_ , and MAM_A denote uniform upper bounds on ‖WB‖2,‖WC‖2,‖WΔ‖2, \|W_B \|_2, \|W_C \|_2, \|W_ \|_2, and ‖A‖2 \|A \|_2, respectively, over the R-local neighborhoods R(Wr)N_R(W^r) and define: MR:=(1+MB)(1+MC)(1+MΔ)(1+MΔ+MA).M_R:=(1+M_B)(1+M_C)(1+M_ )(1+M_ +M_A). Under these conditions, the global optimization trajectory Wrr=0T\W^r\_r=0^T satisfies 1T∑r=0T−1‖∇f(Wr)‖F2=⁡(F0ηET+ηMR4Ξ(1−ρ)8). 1T _r=0^T-1E \|∇ f(W^r) \|_F^2=O ( F_0η ET+ η M_R^4 (1-ρ)^8 ). The quantities N,D,Lx,Lℓ,N,D,L_x,L_ , and HℓH_ affect only the implicit constant and do not alter the displayed scaling with the optimization parameters, parameter envelopes, stability, or sequence length. Proof of Corollary D.8: Setting E=1E=1 in Corollary 4.8 eliminates both terms containing E−1E-1, yielding the equality in the Corollary. ∎ Appendix E Discussion on Additional Distributed Learning Frameworks The FedAvg and FedProx results (Theorems B.5 and B.6) relate convergence of selective SSM training to recurrent stability, parameter norms, local client drift, and data heterogeneity. In this section, we use these quantities to interpret additional federated learning methods, identifying how their local updates and server aggregation interact with the SSM-specific bounds. While we do not establish convergence guarantees for these methods, the discussion provides a qualitative, theory-guided characterization of their convergence behavior, offering intuition and motivating expectations about their performance under different training conditions. The discussion below assumes SGD-style updates, containment in the stated neighborhood, and regularity assumptions. The connection to the AdamW-based Mamba2 experiments is therefore qualitative rather than a direct application of the convergence results. SCAFFOLD. SCAFFOLD [Karimireddy et al., 2020] introduces server and client control variates (correction terms), which act as drift-correction terms for local client updates. In its classical SGD form, the local update is Wkr,e+1=Wkr,e−η⁡(gk(Wkr,e,ξkr,e)+cr−ckr),W_k^r,e+1=W_k^r,e-η (g_k(W_k^r,e; _k^r,e)+c^r-c_k^r ), (434) where cr=∑kpkckrc^r= _kp_kc_k^r approximates the global update direction and ckrc_k^r approximates the client-specific drift direction. More generally, allowing a correction scale λc>0 _c>0, define the corrected local direction dkr,e:=gk(Wkr,e,ξkr,e)+λc(cr−ckr).d_k^r,e:=g_k(W_k^r,e; _k^r,e)+ _c(c^r-c_k^r). (435) The correction term cr−ckrc^r-c_k^r is intended to reduce the component of the local update that points toward the client-specific objective rather than the global objective. Assume that the correction term satisfies: ‖cr−ckr‖F≤CR,\|c^r-c_k^r\|_F≤ C_R, (436) then the corrected direction satisfies: ‖dkr,e‖F≤GSSM,R+λcCR:=GSCAF,R,\|d_k^r,e\|_F≤ G_SSM,R+ _cC_R:=G_SCAF,R, (437) and the corresponding local drift radius can be bounded by RSCAF≤ηEGSCAF,R.R_SCAF≤η EG_SCAF,R. (438) While this mechanism does not directly change the SSM smoothness constant LSSM,RL_SSM,R, and could even increase the worst-case local drift radius through the additional correction term, the main effect of SCAFFOLD is to reduce the component of each client update that is biased toward the local objective rather than the global objective. In other words, SCAFFOLD can be interpreted as modifying the effective heterogeneity term. Specifically, define the corrected heterogeneity measure ΓSCAF,R2:=supW∈R(Wr)∑k=1Kpk‖∇ℒk(W)−∇f(W)+λc(cr−ckr)‖F2. _SCAF,R^2:= _W _R(W^r) _k=1^Kp_k \| _k(W)-∇ f(W)+ _c(c^r-c_k^r) \|_F^2. (439) When the control terms accurately estimate the client-specific drift directions, the correction λc(cr−ckr) _c(c^r-c_k^r) partially cancels the difference ∇ℒk(W)−∇f(W) _k(W)-∇ f(W). In this case, one expects ΓSCAF,R2<ΓSSM,R2. _SCAF,R^2< _SSM,R^2. (440) This motivates a potentially reduced client disagreement, although a convergence bound must also account for the evolution and accuracy of the control variates. This does not imply that SCAFFOLD has a smaller drift radius than FedAvg. In fact, in nearly IID settings, the uncorrected client gradients are already approximately aligned with the global update direction. As a result, the reduction from ΓSSM,R2 _SSM,R^2 to ΓSCAF,R2 _SCAF,R^2 may be limited, while the additional correction term can increase the effective update radius and potentially degrade performance. In contrast, under non-IID data, the heterogeneity correction can provide a meaningful benefit by reducing the statistical drift term, which is amplified in the SSM-specific bound through the architecture-dependent gradient and smoothness quantities. FedAdam. FedAdam [Reddi et al., 2021] keeps the same FedAvg-style local client training phase, but replaces the direct server averaging step with an adaptive server optimizer. Let W¯r+1=∑k=1KpkWkr,E W^r+1= _k=1^Kp_kW_k^r,E (441) denote the standard FedAvg aggregated parameters, and define the server pseudo-gradient: Δr:=W¯r+1−Wr. ^r:= W^r+1-W^r. (442) FedAdam tracks moving averages of the server pseudo-gradient and its coordinatewise square as: mr m^r =β1mr−1+(1−β1)Δr, = _1m^r-1+(1- _1) ^r, (443) vr v^r =β2vr−1+(1−β2)(Δr⊙Δr), = _2v^r-1+(1- _2)( ^r ^r), (444) and updates the server model coordinatewise as Wr+1=Wr+αsmrvr+τs,W^r+1=W^r+ _s m^r v^r+ _s, (445) where αs _s is the server learning rate and τs>0 _s>0 is the denominator offset. Since the client-side training phase is unchanged from FedAvg, FedAdam does not directly reduce the local SSM drift radius. The local client trajectories remain controlled by the same worst-case bound Rlocal≤ηEGSSM,R.R_local≤η EG_SSM,R. (446) The main change is instead introduced at the server: since vr+τs≥τs v^r+ _s≥ _s, the server update satisfies the coarse bound ‖Wr+1−Wr‖F≤αsτs‖mr‖F.\|W^r+1-W^r\|_F≤ _s _s\|m^r\|_F. (447) As such, even though the adaptive server mechanism can be beneficial when it rescales poorly conditioned coordinates in a useful way, it can also increase the effective global update radius when αsτs _s _s is large or when the moment estimates are not well aligned with a stable descent direction. This can be particularly detrimental for selective SSMs - an aggressive server-side adaptive step can increase the effective neighborhood radius, resulting in larger potential values of LSSM,RL_SSM,R and GSSM,RG_SSM,R, and making the global trajectory more sensitive to non-IID client updates. Compared to FedProx, FedAdam does not explicitly constrain local client drift, and compared to SCAFFOLD, FedAdam does not correct the client heterogeneity term. Instead, it rescales the aggregated update after local drift has already occurred. FedAlign. The FedAlign baseline we consider is an adaptation of the state-alignment method proposed in [Keçeci et al., 2025]. The original FedAlign algorithm was developed for federated system identification with classical linear state-space models of the form x~k+1=A~x~k+B~uk,y~k=C~x~k+D~uk x_k+1= A x_k+ Bu_k, y_k= C x_k+ Du_k (448) where the system matrices A~,B~,C~,D~ A, B, C, D are fixed model parameters. In that setting, the main alignment issue is that different clients may learn dynamically equivalent SSMs in different state-coordinate bases, so directly averaging their parameters can distort the global system dynamics. FedAlign addresses this by using similarity transformations to map local SSMs into a common parameter basin before aggregation. Our setting differs from this original SYSID formulation because Mamba-style selective SSMs do not have a single fixed set of state-space matrices applied uniformly across the sequence. Instead, the effective transition and input terms are coupled to learned projections and input-dependent discretization quantities. Therefore, our FedAlign implementation adapts the same state-alignment principle to Mamba2 by aligning coupled SSM state, group, and head-indexed parameter blocks before aggregation. As such, we can write the adapted FedAlign aggregation as Wr+1=∑k=1KpkAlignk⁡(Wkr,E),W^r+1= _k=1^Kp_kAlign_k(W_k^r,E), (449) where AlignkAlign_k is the client-specific alignment map applied to the coupled Mamba2 SSM blocks. Next, as long as AlignkAlign_k is a norm-preserving permutation or matching of coupled coordinates, then ‖Alignk⁡(Wkr,E)−Alignk⁡(Wr)‖F=‖Wkr,E−Wr‖F. \|Align_k(W_k^r,E)-Align_k(W^r) \|_F= \|W_k^r,E-W^r \|_F. (450) The identity preserves distance when the same alignment is applied to both models, not distance to the fixed broadcast model WrW^r. FedAlign can reduce aggregation error from mismatched state coordinates, but norm preservation alone does not control local drift, statistical heterogeneity, or stability-sensitive gradient growth in selective SSMs. FedAvg-like behavior is therefore expected when coordinate mismatch is limited. Containment of aligned models requires separate verification. FedDyn. FedDyn [Acar et al., 2021] combines proximal regularization with a dynamic gradient correction. In the full-participation, weighted formulation, client k maintains a state νkr _k^r, initialized at zero, and optimizes gkr(W):=ℒk(W)−⟨νkr,W⟩F+α2‖W−Wr‖F2,νkr+1:=νkr−α⁡(Wkr,E−Wr),g_k^r(W):=L_k(W)- _k^r,W _F+ α2\|W-W^r\|_F^2, _k^r+1:= _k^r-α(W_k^r,E-W^r), (451) where α>0α>0. Writing νr:=∑kpkνkrν^r:= _kp_k _k^r and W¯r+1:=∑kpkWkr,E W^r+1:= _kp_kW_k^r,E, the server updates are Wr+1=W¯r+1−νr+1/αW^r+1= W^r+1-ν^r+1/α. For local SGD, under the stable-neighborhood conditions of Theorem B.6, given the bound ‖νkr‖F≤Cν\| _k^r\|_F≤ C_ν, and with the learning rate satisfying 0<ηα≤10<ηα≤ 1, the proximal displacement recursion gives: maxe∈0,…,E⁡‖Wkr,e−Wr‖F≤Rdyn:=min⁡R,GSSM,R+Cνα[1−(1−ηα)E]. _e∈\0,…,E\ \|W_k^r,e-W^r \|_F≤ R_dyn:= \R, G_SSM,R+C_να [1-(1-ηα)^E ] \. (452) The geometric factor therefore damps accumulated local displacement. While this bound includes CνC_ν, it ignores cancellation between the local gradient and the dual correction. Accurate tracking can on the other hand produce smaller corrected updates than the inequality suggests, and given a tighter local neighborhood, the SSM parameter envelopes could be reduced. Furthermore, note that the client-gradient disagreement becomes: Γdyn,R2:=supW∈R(Wr)∑kpk‖∇ℒk(W)−∇f(W)−(νkr−νr)‖F2. _dyn,R^2:= _W _R(W^r) _kp_k \| _k(W)-∇ f(W)-( _k^r-ν^r) \|_F^2. (453) As a result, when the centered states accurately track the client-gradient differences, this residual can be smaller than ΓSSM,R2 _SSM,R^2. Therefore, FedDyn combines local-displacement damping with a mechanism for reducing the disagreement highlighted by the FedAvg bound. However, these benefits compete with the additional local smoothness term α and the server displacement: ‖Wr+1−Wr‖F≤Rdyn+‖νr+1‖Fα.\|W^r+1-W^r\|_F≤ R_dyn+ \|ν^r+1\|_Fα. (454) As such, suitable regularization and accurate tracking can make the damping and disagreement beneficial while preserving stable containment, if these improvements outweigh the smoothness constant increase provided the server update also remains within the stable neighborhood. As such, for a well-tuned FedDyn, we expect it to improve on FedAvg, with the benefits being especially highlighted under heterogeneous data conditions through disagreement correction. FedSAM. FedSAM [Qu et al., 2022] retains FedAvg aggregation but uses a sharpness-aware local SGD step with two gradient evaluations: ukr,e:=ℜSAMgk(Wkr,e,ξkr,e)‖gk(Wkr,e,ξkr,e)‖F,Wkr,e+1=Wkr,e−ηgk(Wkr,e+ukr,e,ξkr,e),u_k^r,e:= R_SAM g_k(W_k^r,e; _k^r,e)\|g_k(W_k^r,e; _k^r,e)\|_F, W_k^r,e+1=W_k^r,e-η\,g_k(W_k^r,e+u_k^r,e; _k^r,e), (455) where ℜSAM>0 R_SAM>0 is the perturbation radius, with ukr,e=0u_k^r,e=0 when the first gradient vanishes. Suppose the local iterates lie in R(Wr)N_R(W^r) and the SSM assumptions hold uniformly on R+ℜSAM(Wr)⊆γ,δN_R+ R_SAM(W^r) _γ,δ. The expanded neighborhood contains the perturbed evaluation points, resulting in the local-displacement bound RSAM:=maxe∈0,…,E⁡‖Wkr,e−Wr‖F≤ηEGSSM,R+ℜSAM.R_SAM:= _e∈\0,…,E\\|W_k^r,e-W^r\|_F≤η EG_SSM,R+ R_SAM. (456) The sequence-wise smoothness argument in the proof of Lemma 4.7 also gives ‖gk(Wkr,e+ukr,e,ξkr,e)−gk(Wkr,e,ξkr,e)‖F≤LSSM,R+ℜSAMℜSAM. \|g_k(W_k^r,e+u_k^r,e; _k^r,e)-g_k(W_k^r,e; _k^r,e) \|_F≤ L_SSM,R+ R_SAM R_SAM. (457) As a result, relative to an SGD step from the same iterate and sample, SAM changes the update by at most ηLSSM,R+ℜSAMℜSAMη L_SSM,R+ R_SAM R_SAM. When the perturbation changes the SSM constants only slightly and its accumulated effect over training remains small, we expect FedSAM to perform similarly compared to FedAvg under both nearly IID and heterogeneous data. Note that this expectation relies on the small-perturbation approximation, and possible generalization benefits from flatter solutions are not quantified by these worst-case bounds. FedExP. FedExP [Jhunjhunwala et al., 2023] retains the FedAvg local updates but rescales the aggregated update: Wr+1 W^r+1 =Wr+ηgr∑k=1Kpk(Wkr,E−Wr), =W^r+ _g^r _k=1^Kp_k(W_k^r,E-W^r), (458) ηgr _g^r :=max⁡1,∑k=1Kpk‖Wkr,E−Wr‖F22‖∑k=1Kpk(Wkr,E−Wr)‖F2+ϵg, := \1, _k=1^Kp_k\|W_k^r,E-W^r\|_F^22 \| _k=1^Kp_k(W_k^r,E-W^r) \|_F^2+ _g \, where ϵg>0 _g>0 stabilizes the denominator. By the weighted variance identity, the server scale factor ηgr _g^r exceeds 11 when ∑k=1Kpk‖Wkr,E−∑j=1KpjWjr,E‖F2>‖∑k=1Kpk(Wkr,E−Wr)‖F2+ϵg. _k=1^Kp_k \|W_k^r,E- _j=1^Kp_jW_j^r,E \|_F^2> \| _k=1^Kp_k(W_k^r,E-W^r) \|_F^2+ _g. (459) As a result, rescaling is activated when disagreement among client updates is sufficiently large relative to the squared norm of their average. Otherwise, ηgr=1 _g^r=1 and the server update reduces to FedAvg. For local SGD under the same stable-neighborhood assumptions, neighborhood radius remains Rlocal≤ηEGSSM,R_local≤η EG_SSM,R, however, the server step satisfies: ‖Wr+1−Wr‖F=ηgr‖∑k=1Kpk(Wkr,E−Wr)‖F≤ηgrηEGSSM,R,\|W^r+1-W^r\|_F= _g^r \| _k=1^Kp_k(W_k^r,E-W^r) \|_F≤ _g^rη EG_SSM,R, (460) Since ηgr≥1 _g^r≥ 1, this server-step upper bound is no smaller than the corresponding FedAvg bound, although the actual server displacement need not exceed the local radius. Therefore, server scaling may require a larger neighborhood to cover the server update, with larger parameter envelopes entering the SSM gradient and smoothness bounds. Containment of this expanded neighborhood in γ,δS_γ,δ is further not guaranteed. The frequency of rescaling depends on disagreement among the client updates, which can reflect both data heterogeneity and stochastic sampling. For IID data, we expect sufficiently similar client updates for the server scaling to remain mostly inactive, resulting in FedAvg-like behavior. However, given heterogeneous data, greater dispersion relative to the average update can activate scaling more often. The resulting server steps also amplify any local-drift error in the averaged update, without correcting client disagreement. For the considered SSM setting, we therefore expect worse performance than FedAvg when these displacement and stability costs outweigh the acceleration benefit. This is a qualitative expectation, not a convergence comparison. DiLoCo-style. The original DiLoCo framework [Douillard et al., 2024] combines AdamW local optimization with a Nesterov momentum server update. Our DiLoCo-style baseline retains the shared local phase and updates the server as follows: mor m_o^r =βomor−1+∑k=1Kpk(Wkr,E−Wr), = _om_o^r-1+ _k=1^Kp_k(W_k^r,E-W^r), (461) Wr+1 W^r+1 =Wr+αo(βomor+∑k=1Kpk(Wkr,E−Wr)), =W^r+ _o ( _om_o^r+ _k=1^Kp_k(W_k^r,E-W^r) ), where αo>0 _o>0, βo∈[0,1) _o∈[0,1), and mo−1=0m_o^-1=0. Setting βo=0 _o=0 and αo=1 _o=1 reduces to FedAvg. For the SGD-based comparison, the same stable-neighborhood conditions with a common gradient bound GSSM,RG_SSM,R across rounds hold. Each averaged client update then has norm at most ηEGSSM,Rη EG_SSM,R. As such, unrolling the server momentum recursion we get ‖mor‖F≤ηEGSSM,R/(1−βo)\|m_o^r\|_F≤η EG_SSM,R/(1- _o), and therefore ‖Wr+1−Wr‖F≤αo1−βoηEGSSM,R.\|W^r+1-W^r\|_F≤ _o1- _oη EG_SSM,R. (462) When αo/(1−βo)>1 _o/(1- _o)>1, this upper bound exceeds the corresponding FedAvg bound. Furthermore, unlike FedExP, momentum can accumulate even when clients agree, because it combines averaged updates across rounds. Persistent update directions can increase server displacement, while changes in direction can leave the momentum contribution misaligned with the current averaged update. The server trajectory may therefore require a larger neighborhood and larger SSM parameter envelopes, with stable containment requiring separate verification. The within-round local update rule remains unchanged, and outer momentum does not directly correct client disagreement. As such, depending on the momentum parameters, the displacement and history effects could outweigh the acceleration benefit of momentum, leading to worse performance than FedAvg, although the effects could potentially balance out. Summary. The FedAvg bound (Theorem B.5) connects convergence to LSSM,RL_SSM,R, GSSM,RG_SSM,R, and ℋTVH_TV through local displacement, stable containment, and client disagreement. Under the stated assumptions, FedProx (Theorem B.6) and FedDyn’s proximal term damp displacement bounds, at a smoothness cost. SCAFFOLD and FedDyn can reduce corrected disagreement when their control states accurately track client-gradient differences, though IID data leave less disagreement to correct. FedAdam, FedExP, and DiLoCo-style retain the local phase but modify server updates through coordinatewise adaptation, disagreement-dependent extrapolation, and momentum, respectively. These updates may require larger stable neighborhoods. FedAlign targets coordinate mismatch through norm-preserving permutations, while FedSAM perturbs gradient evaluation points, so FedAvg-like behavior is expected when coordinate mismatch or accumulated perturbation effects, respectively, are small. These theorem-motivated expectations are collected in Table 2 and compared with observations in Section 5. Appendix F Discussion on Model Merging F.1 Model Merging with SSMs Motivation Section 3 formulates distributed learning as the minimization of the aggregated expected sequence loss in (4) through repeated local adaptation and parameter aggregation. The optimization proceeds over T communication rounds, each consisting of E local SGD steps followed by aggregation of the updated model parameters. This framework also provides a natural formulation of one-shot model merging [Wortsman et al., 2022, Ilharco et al., 2023]. Consider the case of task or domain experts adapted from a common pretrained initialization. This setting shares the central operation studied in our distributed framework: combining parameter updates learned from heterogeneous data. Specifically, let M experts share the same selective SSM architecture, parameterization, and pretrained initialization W0W^0. Expert m adapts to a local dataset mD_m associated with a distribution mP_m, with all experts using the sequence-loss formulation of Section 3. To make the connection explicit, suppose each expert starts from Wm0,0=W0W_m^0,0=W^0 and performs E local SGD steps following the update rule in that section. Denote its final parameters and task vector by Wm:=Wm0,E,m:=Wm−W0.W_m:=W_m^0,E, _m:=W_m-W^0. (463) The independently adapted experts are then combined using a merging operator, Wmerge=Merge⁡(Wmm=1M;W0),W_merge=Merge\! (\W_m\_m=1^M;W^0 ), (464) which may use both the expert parameters and the pretrained reference. If the merging operator is chosen to be weighted parameter averaging, the final model parameters are obtained as Wavg=∑m=1MpmWm=W0+∑m=1Mpmm,pm>0,∑m=1Mpm=1,W_avg= _m=1^Mp_mW_m=W^0+ _m=1^Mp_mT_m, p_m>0, _m=1^Mp_m=1, (465) which reduces exactly to one round of FedAvg with K=MK=M and T=1T=1 when the weights in the objective and aggregation are kept the same. This demonstrates that averaging task vectors around a shared initialization provides a direct connection between one-shot model merging and the distributed optimization setting considered in this work. More general merging operators retain this sequence of independent adaptation followed by aggregation, although the differences in how the expert updates are combined limit the applicability of the current theoretical results. Additionally, practical model merging often involves large pretrained networks and extensive independent fine-tuning. Although our formulation permits pretrained initialization and includes a multilayer extension in Appendix C, applying the results to complete pretrained architectures and expert-specific optimization protocols requires verifying additional assumptions. Applying our federated convergence results to one-shot merging requires a neighborhood around W0W^0 containing the full independent adaptation trajectories, without intermediate synchronization. This may require a large radius R and correspondingly large parameter envelopes, increasing LSSM,RL_SSM,R and GSSM,RG_SSM,R and making the bounds loose or potentially uninformative. The required stability and step-size conditions may also become harder to satisfy. When a common stable neighborhood exists, convex averaging preserves containment within it, but stability alone does not guarantee predictive performance across the contributing domains. Nevertheless, the shared structure of independent adaptation followed by aggregation motivates including model merging as a related setting for interpreting our results. In the remainder of this section, we discuss connections between representative merging methods and our analysis, clarify their applicability and limitations, and highlight potential takeaways for merging selective SSMs. We further provide supplementary model merging results in Appendix H.3. F.2 Model Merging Methods: Connections to Selective SSM Analysis For a merged model parameterized by WmergeW_merge with Wmerge∈R(W0)⊆γ,δW_merge _R(W^0) _γ,δ, the uniform regularity conditions of Lemma 4.7 imply f⁡(Wmerge) f (W_merge ) =f⁡(W0+(Wmerge−W0)) =f (W^0+ (W_merge-W^0 ) ) ≤f⁡(W0)+⟨∇f(W0),Wmerge−W0⟩F+LSSM,R2‖Wmerge−W0‖F2. ≤ f (W^0 )+ ∇ f (W^0 ),W_merge-W^0 _F+ L_SSM,R2 \|W_merge-W^0 \|_F^2. (466) Similarly, for a Q-layer model with Wmerge∈Q(W0)⊆,(Q)W_merge _R_Q(W^0) _ γ, δ^(Q), the assumptions of Lemma C.4 give f⁡(Wmerge)≤f⁡(W0)+⟨∇f(W0),Wmerge−W0⟩F+LSSM,Q(Q)2‖Wmerge−W0‖F2,f (W_merge )≤ f (W^0 )+ ∇ f (W^0 ),W_merge-W^0 _F+ L_SSM,R_Q^(Q)2 \|W_merge-W^0 \|_F^2, (467) where the inner product and norm are taken over the full parameter tuple, and LSSM,Q(Q)=HℓΘQ2+LℓΩQL_SSM,R_Q^(Q)=H_ _Q^2+L_ _Q is the existing multi-layer smoothness constant. These inequalities highlight important factors that determine the loss bound after merging. Firstly, the alignment of Wmerge−W0W_merge-W^0 with −∇f(W0)-∇ f(W^0) indicates whether the update initially moves toward a lower aggregated loss. In addition, the merging operator determines the direction and magnitude of this update - different merging methods might rescale or modify the model parameters during aggregation. Finally, the SSM smoothness constant bounds the error in this local loss prediction and depends on the neighborhood over which it is evaluated. The radius R therefore plays two roles: it specifies the maximum displacement of the model parameters ‖Wmerge−W0‖F≤R \|W_merge-W^0 \|_F≤ R (or ‖Wmerge(q)−W0,(q)‖F≤Rq∀q∈1,…,Q \|W_merge^(q)-W^0,(q) \|_F≤ R_q∀ q∈\1,…,Q\) covered by the analysis, and determines the parameter range over which the smoothness bound holds. A larger radius can require larger parameter envelopes MXM_X, leading to larger smoothness constants LSSM,RL_SSM,R (or LSSM,Q(Q)L_SSM,R_Q^(Q)) and a looser loss bound. This motivates examining how different merging methods affect the merged model’s displacement from W0W^0 and the neighborhood needed to contain it. We discuss several representative methods below. Weighted parameter averaging. If all the experts lie in a common stable neighborhood R(W0)N_R(W^0), their weighted average satisfies ‖Wavg−W0‖F≤∑m=1Mpm‖Wm−W0‖F≤R. \|W_avg-W^0 \|_F≤ _m=1^Mp_m \|W_m-W^0 \|_F≤ R. (468) So an averaging merge operation does not enlarge the neighborhood containing the experts, and the same local smoothness bound applies. This property also holds for the multi-layer product neighborhood. Furthermore, Lemma B.4 bounds the weighted squared gradient disagreement at a common parameter tuple by 4GSSM,R2ℋTV4G_SSM,R^2H_TV, showing how data heterogeneity and SSM sensitivity jointly determine this bound. Task Arithmetic. Task Arithmetic [Ilharco et al., 2023] scales combined task vectors. In a weighted, normalized form, WTA=W0+λ¯,¯:=∑m=1Mpmm,λ≥0.W_TA=W^0+λ T, T:= _m=1^Mp_mT_m, λ≥ 0. (469) Here λ=1λ=1 gives parameter averaging, while 0≤λ≤10≤λ≤ 1 preserves containment in the common neighborhood. Larger scales require a separate containment check. In (466), the directional term scales with λ and the quadratic remainder with λ2λ^2. Thus, scaling controls displacement relative to SSM sensitivity, but shrinking the update can also suppress useful adaptation and cannot recover directions lost through cancellation. TIES-Merging. TIES-Merging [Yadav et al., 2023] trims small task-vector entries, elects each coordinate’s sign according to the total retained movement, and averages the nonzero entries agreeing with that sign. The resulting vector is scaled and added to W0W^0. Because coordinates may select different expert subsets, the result need not remain in the experts’ common neighborhood. Moreover, raw magnitude and sign do not account for how changes to A,WΔ,WB,A,W_ ,W_B, and WCW_C interact through the SSM recurrence. Our bounds therefore motivate sensitivity-informed trimming and conflict resolution, but do not establish that TIES improves stability or reduces the merged model’s loss. Practical Takeaways and Limitations. Overall, the model merging discussion motivates jointly considering expert adaptation, merge strength, and SSM sensitivity. The dependence of the bounds on neighborhood radii, recurrent stability, and interactions between layers motivates considering SSM-specific properties when evaluating and designing merging methods. Additionally, the derived relationships remain conditional on the stated assumptions - the local loss bounds require uniform regularity over a common convex stable neighborhood, while applying the federated convergence results additionally requires containment of the full adaptation trajectories. Extensive fine-tuning may therefore require large parameter envelopes and yield loose bounds. These limitations motivate extending the analysis to complete pretrained architectures, practical optimizers, and general merging operators as a promising direction for future work. We complement this discussion with supplementary model-merging experiments and an analysis of their results in Appendix H.3. Appendix G Detailed Experimental Setup In this section, we describe the experimental setup in detail. The experiments fall into three groups: (i) synthetic dataset validation of the derived theoretical bounds, (i) federated language modeling with practical selective state-space models on RedPajama, (i) one-shot model merging supplementary experiments with pretrained Large Mamba2-based models. We detail the datasets, models, as well as training and validation procedures needed for reproducibility. G.1 Datasets Synthetic Validation. The synthetic datasets are autoregressive token sequences sampled from frozen teacher selective SSMs that implement the recurrence of Section 3 exactly, so that every assumption of Theorem 4.4 holds for the data-generating model (rather than being assumed to hold). A teacher consists only of the four parameter blocks W=(WΔ,WB,WC,A)W= (W_ ,W_B,W_C,A ) of a layer (no embedding, convolution, residual path, normalization, MLP, or Transformer components). The inputs are one-hot encoded: xtx_t is one of the standard basis vectors e1,…,eDe_1,…,e_D of ℝDR^D for every t. This formulation makes the input dimension D also act as the vocabulary size (since there are D different basis vectors). Since ‖xt‖2=1\|x_t\|_2=1, Assumption 4.3 holds with Lx=1L_x=1. The teacher SSM is turned into a generative model as follows. For a partially generated sequence x1,…,xtx_1,…,x_t, running it through the recurrence, starting from h0=N×Dh_0=0_N× D, produces an output yt∈ℝDy_t ^D, which has one coordinate per token in the vocabulary. We therefore treat the d-th coordinate yt,dy_t,d as the logit of the event that the next input is token d, i.e., that xt+1=edx_t+1=e_d. We then take a softmax with temperature τ>0τ>0 to convert these D logits into a probability vector over the tokens, and draw the next input as follows: Pr[xt+1=ed|x1,…,xt]=[softmax(yt/τ)]d=exp⁡(yt,d/τ)∑d′=1Dexp⁡(yt,d′/τ),d=1,…,D. [x_t+1=e_d\, |\,x_1,…,x_t ]= [softmax (y_t/τ ) ]_d= (y_t,d/τ ) _d =1^D (y_t,d /τ ), d=1,…,D. Table 3 summarizes the single-layer dataset generation over different settings. Table 3: Generation settings of the synthetic teacher datasets. Grid entries are searched per configuration in the fixed order given in the text. Setting Value Input/vocabulary dimension D 4,8,16\4,8,16\ State dimension N 4,8,16\4,8,16\ Stability limit ρ 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 S 6464 stored tokens (1616 burn-in tokens discarded, 8080 generated) Training / test sequences 4,0964,096 / 1,0241,024 Sampling seeds (train / test) 1,0001,000 / 1,001,0031,001,003 First input x1x_1 uniform over e1,…,eD\e_1,…,e_D\ Teacher seeds s 0,1,2\0,1,2\ (grid) Gain g, temperature τ grid WΔW_ , WCW_C i.i.d. ⁡(0,0.752)N(0,0.75^2) entries WBW_B i.i.d. ⁡(0,(0.75g)2)N(0,(0.75\,g)^2) entries A An=−γ⁡(1+0.25|zn|)A_n=-γ(1+0.25|z_n|), zn∼⁡(0,1)z_n (0,1) Step margin δ ⌊105δ^⌋/105 10^5 δ /10^5 (realized minimum step of WΔW_ ) Transition margin γ −ln(ρ)/δ- (ρ)/δ Parameter precision single precision (float32) Single-layer synthetic datasets. To generate the single-layer datasets, we instantiate one frozen teacher SSM for every combination of the input dimension D∈4,8,16D∈\4,8,16\, the state dimension N∈4,8,16N∈\4,8,16\, and the stability limit ρ∈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\ (from Definition 4.1), which gives 3×3×8=723× 3× 8=72 different frozen teachers. Each such teacher, with its specific combination of D, N, and ρ, then generates 4,0964,096 training and 1,0241,024 test sequences of length S=64S=64 according to the procedure described above. Teacher initialization and stability targeting. For a teacher SSM corresponding to a combination of D, N, and ρ, its weights are drawn with i.i.d. Gaussian entries as follows: (WB)n,d∼⁡(0,(0.75g)2) (W_B )_n,d (0,(0.75g)^2 ), (WC)n,d∼⁡(0,0.752) (W_C )_n,d (0,0.75^2 ), (WΔ)d,d′∼⁡(0,0.752) (W_ )_d,d (0,0.75^2 ), where g is a gain multiplier on the input projection described below. The margins (γ,δ)(γ,δ) are chosen such that the teacher belongs to γ,δS_γ,δ with ρ=exp⁡(−γδ)ρ= (-γδ) equal to the specified value while satisfying both constraints of Definition 4.1. Since the inputs are one-hot, WΔxtW_ x_t is simply a column of WΔW_ , so the time-scale constraint of Definition 4.1 reduces to the entry-wise condition Softplus⁡((WΔ)d,d′)≥δSoftplus ( (W_ )_d,d )≥δ for all d,d′d,d . We therefore draw WΔW_ first and set δ to the smallest step it produces, δ^:=mind,d′⁡Softplus⁡((WΔ)d,d′)=Softplus⁡(mind,d′⁡(WΔ)d,d′),δ:=⌊105δ^⌋/105, δ:= _d,d Softplus ( (W_ )_d,d )=Softplus ( _d,d (W_ )_d,d ), δ:= 10^5\, δ /10^5, rounded down so that δ≤δ^δ≤ δ holds exactly, and then γ:=−ln(ρ)/δγ:=- (ρ)/δ for the designated ρ. The parameters of the transition block are then generated as follows: An=−γ(1+0.25|zn|),zn∼(0,1),n=1,…,N,A_n=-γ (1+0.25\,|z_n| ), z_n (0,1), n=1,…,N, so that An≤−γA_n≤-γ holds with a multiplicative margin. With this construction, the realized stability limit ρ^:=maxn,d,d′⁡(A¯t(ed′))n,d=exp⁡(−γ^δ^) ρ:= _n,d,d ( A_t(e_d ) )_n,d= (- γ δ) with γ^:=minn⁡(−An) γ:= _n(-A_n) satisfies ρ^=ρ(δ^/δ)(1+0.25minn|zn|) ρ=ρ^( δ/δ)(1+0.25 _n|z_n|) and tracks ρ, including as ρ→1ρ→ 1. Across all generated datasets ρ−ρ^∈[0,0.03]ρ- ρ∈[0,0.03]. By construction, An≤−γA_n≤-γ, Softplus⁡((WΔ)d,d′)≥δSoftplus ((W_ )_d,d )≥δ, and ρ^<1 ρ<1 hold. Since the admissible inputs form the finite set e1,…,eD\e_1,…,e_D\, these conditions are also verified exhaustively over all D one-hot inputs before a teacher is accepted. A clamping projection onto γ,δS_γ,δ is available as a safeguard, but it was not active for any of the 7272 accepted teachers. Sampling. Given the frozen teacher with a combination of D, N, and ρ and the frozen weights constructed as described above, a sequence is generated by drawing x1x_1 uniformly from e1,…,eD\e_1,…,e_D\, starting from h0=N×Dh_0=0_N× D, and sampling xt+1x_t+1 as described above for 7979 further steps. The first 1616 tokens are discarded as burn-in and the remaining S=64S=64 tokens are stored as a single sequence. The burn-in aims to reduce the stored prefix’s dependence on the first input. The recurrent state at the start of the stored window is not persisted, and every subsequent teacher-forced or learner evaluation restarts from h0=N×Dh_0=0_N× D at the first stored token, which is the convention of the theory. The training and test splits are sampled with separate fixed seeds (1,0001,000 and 1,001,0031,001,003), reused in every configuration. Each dataset stores the token indices as 8-bit integers together with the frozen teacher parameters, the generation settings, and the stability certificate, and is never modified afterwards. The learner only sees the stored tokens and never accesses the teacher. Acceptance criteria. A generated candidate dataset is kept only if it passes the fixed criteria of Table 4, which exclude degenerate teachers (whose sequences collapse into a dominant token, or become deterministic, constant, or uninformative), while ensuring that the next-token task carries predictive information. All diagnostics of the table are computed on the union of the training and test splits, unless stated otherwise. The central quantity is the teacher-forced negative log-likelihood (NLL): the stored tokens of a sequence are fed to the teacher as inputs, starting from h0=N×Dh_0=0_N× D, instead of re-sampling from it, and at every position t=1,…,S−1t=1,…,S-1 we record −ln(Pr[xt+1∣x1,…,xt])- ( [x_t+1 x_1,…,x_t] ), i.e., minus the log-probability that the teacher (at the sampling temperature τ) assigns to the stored next token. The NLL is the average of these values over all positions and sequences and measures how predictable the generated data are for the very model that produced them. A model that assigns probability 1/D1/D to every token has NLL ln⁡D D, which serves as the uniform baseline. The informativeness criterion therefore requires the sequences to be predictable beyond chance, while the perplexity criterion exp⁡(NLL)≥1.20 (NLL)≥ 1.20 excludes near-deterministic teachers. The mutual information is estimated from the empirical frequencies of consecutive token pairs. To avoid generating full datasets for clearly degenerate candidates, each candidate is first screened on a small draw of 256256 training and 128128 test sequences, and only if this screen passes is the full dataset generated and checked again against the same thresholds, which are never relaxed. 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 D tokens over all stored positions, ignoring position and context. NLL denotes the teacher-forced negative log-likelihood at the sampling temperature τ with h0=h_0=0 per stored sequence. Criterion Threshold Global token balance largest probability in the token marginal ≤1/D+0.10≤ 1/D+0.10 Marginal entropy Shannon entropy of the token marginal ≥0.90ln⁡D≥ 0.90 D (the maximum ln⁡D D is attained by the uniform distribution) Per-sequence majority (test split) for each test sequence, the fraction of its S positions occupied by its most frequent token; the 9595th percentile of this fraction over the test sequences ≤1/D+0.20≤ 1/D+0.20 Class coverage each of the D tokens occurs at least once in the training split and at least once in the test split Constant sequences fraction of sequences whose S tokens are all identical ≤0.01≤ 0.01 Sequence diversity number of distinct sequences divided by the number of sequences ≥0.95≥ 0.95 Teacher informativeness ln⁡D−NLL≥0.02 D-NLL≥ 0.02, i.e., the teacher beats the uniform predictor No deterministic collapse teacher perplexity exp⁡(NLL)≥1.20 (NLL)≥ 1.20 Predictive information mutual information between consecutive tokens xtx_t and xt+1x_t+1, estimated from their empirical joint frequencies, ≥0.005≥ 0.005 Finiteness all diagnostics, logits, steps Δt _t, transitions A¯t A_t, and hidden states are finite (no overflow) Stability the teacher is certified to lie in γ,δS_γ,δ with realized stability limit ρ^<1 ρ<1 Gain and temperature choices. As ρ→1ρ→ 1, the hidden state of a random teacher accumulates input contributions over a horizon of order 1/(1−ρ)1/(1-ρ), so its logits grow and sampling collapses into a few tokens, which violates the balance and entropy criteria. Instead of relaxing the criteria, we search a fixed grid of input-projection gains and sampling temperatures, (g,τ)∈ (g,τ)∈ \ (1,0.7),(1,1),(0.5,0.7),(0.5,1),(0.25,0.7),(0.25,1),(0.15,0.7), (1,0.7),(1,1),(0.5,0.7),(0.5,1),(0.25,0.7),(0.25,1),(0.15,0.7), (0.15,1),(0.1,0.7),(0.1,1),(0.05,0.7),(0.05,1),(0.03,1), (0.15,1),(0.1,0.7),(0.1,1),(0.05,0.7),(0.05,1),(0.03,1) \, in this order for the teacher seeds s∈0,1,2s∈\0,1,2\ (outer loop), and accept the first candidate that passes every criterion in Table 4. Reducing g scales WBW_B (and therefore B¯t B_t) and the accumulated state without touching WΔW_ or A, so the margins (γ,δ)(γ,δ) and ρ remain unchanged. Over the 7272 dataset configurations, 1818 were accepted at the first attempt and the longest search took 3232 attempts. Teacher seed 00 was accepted in 6363 configurations, seed 11 in 55, and seed 22 in 44. The accepted gains were g=1g=1 (3030 configurations), 0.50.5 (99), 0.250.25 (1515), 0.150.15 (88), 0.10.1 (22), 0.050.05 (66), and 0.030.03 (22). A gain below one was required in every configuration with ρ≥0.99ρ≥ 0.99 and in only 77 of the 2727 configurations with ρ≤0.7ρ≤ 0.7. The accepted temperatures were τ=0.7τ=0.7 (3939 configurations) and τ=1τ=1 (3333 configurations). The resulting teachers improve on the uniform NLL by 0.020.02 to 0.970.97, have a consecutive-token mutual information between 0.0080.008 and 0.900.90 and a teacher perplexity between 1.51.5 and 15.715.7. The most frequent token has frequency at most 1.76/D1.76/D, the normalized marginal entropy is at least 0.940.94, and at least 98%98\% of the stored sequences are unique. Federated language modeling (RedPajama). The federated experiments use the SlimPajama-627B [Together Computer, 2023] release of RedPajama, available as cerebras/SlimPajama-627B.. Of the seven SlimPajama domains: arxiv, book, c4, common_crawl (c), github, stackexchange, wikipedia, we use four-domain combinations to instantiate four federated clients: • 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. Due to copyright restrictions, the book domain is no longer accessible, and therefore not used in our experiments. For every client, we set the maximum number of sequences to sequences_per_client = 10,000. Each client’s dataset is split into 80% training and 20% validation data using train_val_split = 0.8. Due to the limited number of available sequences in certain domains, the effective per-client budgets are reduced for some settings. The resulting dataset statistics are summarized in Table 5. Table 5: Effective client train/validation budgets and server validation sizes across heterogeneity settings. Setting β Train cap Val cap Train limiter data Valid limiter data Server val Setting 1 0.0 7,427 1,857 stackexchange stackexchange 7,424 Setting 1 0.7 7,427 1,857 stackexchange stackexchange 7,431 Setting 1 1.0 7,427 1,857 stackexchange stackexchange 7,428 Setting 2 0.0 7,367 1,842 c4 c4 7,360 Setting 2 0.7 7,367 1,842 c4 c4 7,364 Setting 2 1.0 7,367 1,842 c4 c4 7,368 Setting 3 0.0 7,367 1,842 c4 c4 7,360 Setting 3 0.7 7,367 1,842 c4 c4 7,364 Setting 3 1.0 7,367 1,842 c4 c4 7,368 Setting 4 0.0 7,427 1,857 stackexchange stackexchange 7,424 Setting 4 0.7 7,427 1,857 stackexchange stackexchange 7,431 Setting 4 1.0 7,427 1,857 stackexchange stackexchange 7,428 We use β to label three domain-allocation regimes: • β=0.0β=0.0 (IID): each client receives an equal mixture over all K=4K=4 domains. • β=1.0β=1.0 (pure non-IID): each client receives data exclusively from its primary domain. • β=0.7β=0.7 (mixed): each client receives 0.70.7 of its data from its primary domain and 0.10.1 from each of the other three. Concretely, with sequences_per_client = 10,000 and an 80/20 train/validation split, each client typically contains 8,000 training chunks and 2,000 validation chunks. The server held-out evaluation set is the union of all client validation pools, giving approximately 8,000 validation chunks for K=4K=4 clients. For each configuration, we use the Hugging Face gpt2 tokenizer (50,257 entries), with EOS reused for padding. The model configuration sets vocab_size = 50,277 and pad_vocab_size_multiple = 8, yielding 50,280 embedding/output rows. Model Merging. For model merging, we evaluate on four text classification datasets: AG News [Zhang et al., 2015], DBpedia [Lehmann et al., 2015], TREC [Li and Roth, 2002], and IMDB [Maas et al., 2011]. The dataset statistics and evaluation limits are summarized in Table 6. Each task is formulated as causal-LM scoring without a classification head. At evaluation time, we score prompt + ‘ ’ + label for every candidate label and select the label with the highest mean per-token log-probability over the label tokens. We use the mean, rather than the sum, so that labels of different lengths remain comparable. The evaluation pool is sub-sampled with a fixed seed of 99991 so that the evaluation split is identical across training seeds. The training pool is sub-sampled with the per-run training seed (seed ∈0,1,2∈0,1,2) so that the three runs differ both in optimization noise and in their training subsamples. Table 6: Task datasets and sequence limits for causal-LM evaluation. Dataset Task # Classes Train cap Eval cap max_seq_len ag_news News topic 4 40,000 5,000 256 dbpedia_14 Wikipedia topic 14 40,000 5,000 256 trec Question type 6 ≈5,500≈ 5,500 (full) 5,000 256 imdb Sentiment 2 40,000 5,000 256 The prompt templates are defined in Table 7: Table 7: Prompt templates and candidate labels used for causal-LM scoring. Dataset Prompt template and labels ag_news Article: text : Labels: World, Sports, Business, Tech dbpedia_14 <title-aware template>...:Category: Labels: 14-way Wikipedia topic labels trec Question: text : Labels: Description, Entity, Abbreviation, Person, Location, Number imdb Review: text : Labels: negative, positive Tokenization uses each model’s own tokenizer, namely GPT-2 BPE. Prompts are truncated or padded to max_seq_len = 256, with 16 tokens reserved for the label completion and template tail. During tokenization, we construct a boolean label_mask. This mask is used during fine-tuning to compute the loss only over label-token positions, and during evaluation to average scores only over label tokens. All datasets are used under their original licenses and terms of use. G.2 Models Synthetic Validation. The learner SSM has the exact same architecture as the teacher: a single-layer selective SSM as defined in Section 3, with the four parameter blocks W=(WΔ,WB,WC,A)W= (W_ ,W_B,W_C,A ) of sizes D×D× D, N×DN× D, N×DN× D, and N×1N× 1, one-hot inputs, and no embedding, convolution, residual path, normalization, MLP, or Transformer components. The output yt∈ℝDy_t ^D is used directly as the logit vector for next-token prediction. In every configuration, the learner uses the same input dimension D and state dimension N as the teacher that generated its dataset (D must match the vocabulary), so its parameter count is D2+2ND+ND^2+2ND+N (Table 8). The learner never receives the teacher parameters - the only teacher-related quantities it uses are the margins (γ,δ)(γ,δ) of its configuration, which define the stability set γ,δS_γ,δ in which it is trained. Table 8: Learner parameter counts D2+2ND+ND^2+2ND+N for the synthetic configurations. N=4N=4 N=8N=8 N=16N=16 D=4D=4 5252 8888 160160 D=8D=8 132132 200200 336336 D=16D=16 388388 520520 784784 Initialization. For a learner seed s∈0,1,2s∈\0,1,2\, drawn independently of the teacher seed, the projections are initialized with i.i.d. Gaussian entries (WΔ)d,d′,(WB)n,d,(WC)n,d∼⁡(0,0.12) (W_ )_d,d , (W_B )_n,d, (W_C )_n,d (0,0.1^2), and the transition block as An=−(γ+ϵ)−0.1|zn|,zn∼(0,1),n=1,…,N,A_n=- (γ+ε )-0.1\,|z_n|, z_n (0,1), n=1,…,N, with a safety margin ϵ=10−4ε=10^-4. The initial parameters are then projected onto γ+ϵ,δ+ϵS_γ+ε,δ+ε by clamping An≤−(γ+ϵ)A_n≤-(γ+ε) and (WΔ)d,d′≥Softplus−1⁡(δ+ϵ) (W_ )_d,d ^-1(δ+ε), and membership in γ,δS_γ,δ is verified over all D one-hot inputs, exactly as for the teachers. The initialization scale 0.10.1 is much smaller than the teacher’s 0.750.75, so every learner starts close to a uniform next-token predictor. The same projection is applied after every optimizer update during training, so the learner remains in γ,δS_γ,δ with realized transitions bounded by ρ throughout training. All learner parameters are stored and updated in single precision (float32). Federated Training. All federated runs train Mamba-2 selective state-space models from scratch using mamba_ssm ’s MambaLMHeadModel with the configuration: • d_model = 256 • ssm_cfg.layer = Mamba2 – default d_state = 128 – d_conv = 4 – expand = 2 – headdim = 64 – ngroups = 1 • vocab_size = 50,277 • pad_vocab_multiple = 8, giving padded vocabulary size 50,28050,280 • rms_norm = True • residual_in_fp32 = True • fused_add_norm = True • dtype = float32 (with fp16 autocast inside forward pass) • depth ∈4,16∈\4,16\ (observed depths: 4 and 16) i.e. we consider custom Mamba2 models of depth 44 or depth 1616. The parameter counts are summarized in Table 9. Table 9: Parameter breakdown by model depth. Depth Total params SSM (mixer) params Embedding params Per-layer mixer 4 14,733,664 1,860,704 12,871,680 465,176 16 20,318,848 7,442,816 12,871,680 465,176 Notes: the input/output embedding are shared by the LM head (50,280×256=12,871,68050,280× 256=12,871,680 parameters). The SSM (mixer) params denotes the parameter count restricted to tensors matching *.mixer.*, including in_proj, conv1d, A_log, D, dt_bias, the block-gated RMSNorm, and out_proj. 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. Backbone Layout SSM fraction Total params mamba2 4×4× Mamba2 1.001.00 14,733,66414,733,664 hybrid75 attention at 2\2\ 0.750.75 14,531,65614,531,656 hybrid50 attention at 1,3\1,3\ 0.500.50 14,329,64814,329,648 hybrid50alt attention at 0,2\0,2\ 0.500.50 14,329,64814,329,648 hybrid25 attention at 0,2,3\0,2,3\ 0.250.25 14,127,64014,127,640 attn0 4×4× attention 0.000.00 13,925,63213,925,632 Hybrid and Attention Backbones. The backbone ablation (see the ablations in the Practical Distributed Learning Framework subsection) replaces some or all of the four Mamba2 mixers of the depth-44 model by causal multi-head attention layers, keeping d_model =256=256, the tied GPT-2 embedding and LM head, the RMSNorm layers, and the residual structure unchanged, so that the mixer of every layer is either a Mamba2 block with the configuration above or an attention block. Attention layers are mamba_ssm’s MHA (Multi-Head Attention) with 88 heads of dimension 3232, rotary position embeddings over the full head dimension, causal masking, biases on the input (qkv) and output projections, and no feed-forward sublayer. The attention layers are spread evenly through the stack. With nattnn_attn attention layers among the L=4L=4 layers, the j-th attention layer (counting from j=0j=0) is placed at index ⌊(j+0.5)L/nattn⌋ (j+0.5)\,L/n_attn , which puts a single attention layer in the middle of the stack (index 22), two attention layers at indices 11 and 33, and three at indices 00, 22, and 33, with the remaining layers being Mamba2 mixers. To check that the placement itself does not drive the results, hybrid50alt keeps two attention layers but moves them to indices 00 and 22. Table 10 lists the resulting layouts and parameter counts: an attention layer has 263,168263,168 parameters against 465,176465,176 for a Mamba2 mixer, so the total parameter count decreases with the attention fraction. The attention-only model (attn0) serves as the zero-SSM control (since it has no SSM state to align). Model Merging. The model merging experiments use public Mamba2 checkpoints loaded from Hugging Face with MambaLMHeadModel.from_pretrained. The results evaluate the 130M, 370M, and 780M checkpoints. Table 11: Public Mamba2 checkpoints used for model merging. Checkpoint Total params state-spaces/mamba2-130m 130M state-spaces/mamba2-370m 370M state-spaces/mamba2-780m 780M All checkpoints are loaded with dtype = bfloat16. The base state θ0θ^0 is the unmodified pretrained checkpoint and is fixed once at the start of each experiment. The per-task specialists θm _m are produced by task fine-tuning (details in G.5), and saved to the per-experiment directory. The pretrained models and the mamba_ssm module are used under their original licenses and terms of use. G.3 Synthetic Distributed Learning Framework For every configuration (D,N,ρ)(D,N,ρ) and learner seed s∈0,1,2s∈\0,1,2\, the synthetic pipeline runs four stages on the dataset corresponding to that configuration: (i) centralized training of the learner of the Models subsection, (i) evaluation of the gradient bound of Theorem 4.4 at the stored checkpoints, (i) evaluation of the smoothness constants of Lemma 4.7 at the same checkpoints, and (iv) federated training with FedAvg and FedProx, which measures the quantities entering Theorem B.5 and Theorem B.6. Stages (i) and (i) never train, they only read checkpoints. All stages use the margins (γ,δ)(γ,δ) of the configuration. Table 12 lists the shared settings. Table 12: Shared settings of the synthetic centralized and federated runs. Component Value Loss mean softmax cross-entropy over the S−1S-1 next-token positions Optimizer SGD, no momentum, weight decay, or gradient clipping Learning rate η 0.050.05 Batch size 6464 sequences Centralized updates 2,0002,000 (epochs of 6464 updates, reshuffled every epoch) Checkpoints and evaluation every 100100 updates; bound and smoothness checkpoints at updates 00, 500500, 1,0001,000, 1,5001,500, 2,0002,000 Stability projection onto γ+ϵ,δ+ϵS_γ+ε,δ+ε with ϵ=10−4ε=10^-4 and certificate check after every update Learner seeds 0,1,2\0,1,2\ per configuration Precision training in float32; bound, smoothness, and certificate evaluation in float64 Centralized Training. The learner minimizes the mean teacher-forced next-token cross-entropy over the S−1S-1 shifted positions of every training sequence, i.e., the sequence loss of (3) with ℓt _t the softmax cross-entropy and yt∗=xt+1y_t^*=x_t+1. SGD uses no momentum, weight decay, or clipping, while projection and epoch reshuffling differ from the analyzed independent-sampling protocol. Each epoch shuffles the 4,0964,096 training sequences with a seed derived from the learner seed and the epoch index and cuts them into batches of 6464, running for 2,0002,000 updates. After every update the parameters are projected onto γ+ϵ,δ+ϵS_γ+ε,δ+ε as at initialization, and membership in γ,δS_γ,δ is re-verified over all D one-hot inputs. Checkpoints are written every 100100 updates, and the test NLL, perplexity, and next-token accuracy are evaluated on the 1,0241,024 test sequences at the same frequency, always restarting from h0=N×Dh_0=0_N× D at the first stored token. Gradient-Bound Evaluation. At the five checkpoints of Table 12 and at the frozen teacher, every one of the 1,0241,024 test sequences is processed in float64, and the gradient norm of its sequence loss with respect to the four blocks is computed as follows: ‖∇WℒS‖F=(∑X∈Δ,B,C,A‖∇WXℒS‖F2)1/2\| _WL^S\|_F= ( _X∈\ ,B,C,A\\| _W_XL^S\|_F^2 )^1/2. It is compared with two constants evaluated at the same parameters. The configured bound SSM(W)G_SSM(W) is Theorem 4.4 with its analytic constants, the configured (γ,δ,ρ)(γ,δ,ρ), Lx=1L_x=1, and Lℓ=2L_ = 2 (the gradient of the softmax cross-entropy with respect to the logits has norm at most 2 2). The tightened bound SSMtight(W)G^tight_SSM(W) uses realized stability margins, exact one-hot projection bounds, and finite-history sums: the realized margins γ^=minn⁡(−An) γ= _n(-A_n), δ δ, and ρ^=exp⁡(−γ^δ^) ρ= (- γ δ) replace (γ,δ,ρ)(γ,δ,ρ), maxd⁡‖Softplus⁡(WΔed)‖2 _d\|Softplus(W_ e_d)\|_2 replaces KΔK_ , the largest column norms maxd⁡‖WBed‖2 _d\|W_Be_d\|_2 and maxd⁡‖WCed‖2 _d\|W_Ce_d\|_2 replace ‖WB‖2\|W_B\|_2 and ‖WC‖2\|W_C\|_2, and every geometric factor 1/(1−ρ)1/(1-ρ) is truncated to (1−ρ^S−1)/(1−ρ^)(1- ρ^S-1)/(1- ρ), since a stored sequence contributes at most S−1S-1 terms to each history sum. These replacements upper-bound the same one-hot projection and history terms and do not exceed the configured constants. A violation is recorded when the measured norm exceeds the bound by more than a relative tolerance of 10−810^-8. The utilization is the ratio of the measured norm to the bound. The gradient of the batch-mean loss over all test sequences is evaluated as well and checked against the mean of the per-sequence norms (triangle inequality) and against the bound. Smoothness Evaluation. At the same checkpoints, 3232 test sequences drawn once with a fixed seed are processed in float64. For every sequence, (a) the operator norm of the Hessian of the sequence loss is estimated by power iteration with exact Hessian-vector products (at most 1515 iterations, stopping when the Rayleigh quotient changes by less than a relative 10−910^-9) and compared with LSSM,RL_SSM,R at R=0R=0, (b) the finite-difference ratio ‖∇ℒS(W+hU)−∇ℒS(W)‖F/h\| ^S(W+hU)- ^S(W)\|_F/h with h=10−3h=10^-3 is evaluated along three unit directions U (two random and the normalized gradient) and compared with LSSM,hL_SSM,h whenever the ball of radius h is guaranteed to lie in γ,δS_γ,δ. Between consecutive learner checkpoints Wa→WbW_a→ W_b, the per-sequence ratio maxi⁡‖∇ℓi(Wb)−∇ℓi(Wa)‖F/‖Wb−Wa‖F _i\|∇ _i(W_b)-∇ _i(W_a)\|_F/\|W_b-W_a\|_F is compared with LSSM,RL_SSM,R centered at WaW_a with R=‖Wb−Wa‖FR=\|W_b-W_a\|_F, and |SSM(Wb)−SSM(Wa)|/‖Wb−Wa‖F|G_SSM(W_b)-G_SSM(W_a)|/\|W_b-W_a\|_F with LSSM,RL_G_SSM,R of Lemma B.1. The loss-regularity constants are Lℓ=2L_ = 2 and Hℓ=1/2H_ =1/2. Federated Training. FedAvg and FedProx are run on the same learner with K clients holding disjoint, equal-size shards of the 4,0964,096 training sequences and weights pk=1/Kp_k=1/K. Under the IID partition, the sequences are randomly permuted (seeded by the learner seed) and cut into K shards. Under the dominant-token partition, the sequences are first sorted by their most frequent token, then by the fraction of positions that token occupies, and then cut into K shards, so that clients differ systematically in token composition. Every client starts each round from the broadcast model and performs E local SGD updates with the settings of Table 12, each on a fresh batch of 6464 sequences drawn with replacement from the client’s shard, followed by the same projection and verification as in centralized training. FedProx adds μ⁡(W−Wr)μ(W-W^r) to every local gradient, taking projected SGD steps on FkrF_k^r from Section 3. The server forms Wr+1=∑kpkWkr,EW^r+1= _kp_kW_k^r,E in float64 and applies the projection once more. All runs start from the centralized learner initialization of the same seed and use T=60T=60 rounds. Table 13 lists the settings: the base setting is run on all 7272 configurations, and the one-factor variations of K, E, μ, and the partition are run on the 1616 core configurations, namely the ρ axis at D=N=8D=N=8 and all (D,N)(D,N) pairs at ρ=0.95ρ=0.95. With three seeds this gives 1,2001,200 federated runs. 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. Component Value Algorithms FedAvg; FedProx Base setting K=4K=4 clients, E=16E=16 local updates, IID partition Communication rounds T 6060 Client-count variation K∈2,8,32K∈\2,8,32\ Local-step variation E∈1,4,64E∈\1,4,64\ Proximal variation (FedProx) μ∈0.01,1μ∈\0.01,1\ Partition variation dominant-token (non-IID) Evaluation at round 00 and every 55 rounds; ‖∇f(Wr)‖F2\|∇ f(W^r)\|_F^2 at every broadcast Seeds 0,1,2\0,1,2\, shared across the settings of a configuration The following quantities are recorded. At every evaluated round, the subset gradient norm ‖∇f(Wr)‖F\|∇ f(W^r)\|_F on a fixed subset of 1,0241,024 training sequences is compared with the configured and tightened SSM(Wr)G_SSM(W^r). The client drift ‖Wkr,E−Wr‖F\|W_k^r,E-W^r\|_F (mean and maximum over clients) and the largest displacement of any local iterate ‖Wkr,e−Wr‖F\|W_k^r,e-W^r\|_F are compared with the radius R. The heterogeneity ∑kpk‖∇ℒk(Wr)−∇f(Wr)‖F2 _kp_k\| _k(W^r)-∇ f(W^r)\|_F^2, with client gradients over full shards, is compared with 4SSM(Wr)2∑kpkTVk24\,G_SSM(W^r)^2 _kp_k\,TV_k^2 of Lemma B.4, where TVkTV_k is the exact total variation between the empirical distribution of client k over whole sequences and the pooled distribution (unigram and bigram variants are recorded for reference but do not enter the lemma). The test NLL of the global model is evaluated on the 1,0241,024 test sequences. Reference radii and envelopes are set from initialization: R=ηESSM(W0)R=η E\,G_SSM(W^0) and MX=‖WX0‖2+(T+1)RM_X=\|W_X^0\|_2+(T+1)R for each block X, from which GSSM,RG_SSM,R, LSSM,RL_SSM,R, and LSSM,RL_G_SSM,R are computed. For FedProx the radius is Rμ=min⁡R,GSSM,R[1−(1−ημ)E]/μR_μ= \R,\,G_SSM,R[1-(1-ημ)^E]/μ \ with correspondingly reduced envelopes. Beyond the sufficient step-size range, we log iterate containment, projection activity, and the tests ηEL≤1/2η EL≤ 1/2 and 12η2(L+μ)2E(E−1)≤112η^2(L+μ)^2E(E-1)≤ 1. We compare the right-hand side of Theorem B.5 (coefficients 4,4,6,484,4,6,48) or Theorem B.6 (coefficients 3,3,3,243,3,3,24) with the subset estimate of 1T∑r<T‖∇f(Wr)‖F2 1T _r<T\|∇ f(W^r)\|_F^2, using the full-training initial loss f⁡(W0)f(W^0) and f∗=0f^*=0. G.4 Practical Distributed Learning Framework The federated framework uses K=4K=4 clients unless stated otherwise (see the client-count ablation) and initializes all clients from the same randomly initialized model θ0θ^0. At the start of every round r, the server broadcasts the global model θrθ^r, each client k runs E local optimization steps on its private data starting from θrθ^r, and returns its updated parameters θkr,E _k^r,E to a server module that performs algorithm-specific aggregation. After aggregation, the new global model θr+1θ^r+1 is evaluated on the union of all clients’ local validation data (server validation pool). The per-round per-client training perplexities and the server validation perplexity are saved. Shared Training Schedule. All algorithms are evaluated under the same overall optimization schedule. Across communication rounds, we keep the number of clients, participation rate, local optimizer, batch construction, sequence length, and random seeds fixed. The only differences across methods are the algorithm-specific server or client update rule and the corresponding method-specific hyperparameters. The shared hyperparameters are summarized in Table 14. Local training runs in single precision with mixed-precision (fp16) autocast and dynamic loss scaling, and gradients are clipped after unscaling. FedSAM is the one exception and runs in full float32. A fresh AdamW optimizer is created at the start of every round, so no optimizer state carries across rounds. Table 14: Shared training hyperparameters used across all federated algorithms. Component Value Number of clients, K 44 Participation rate 1.01.0; all clients participate each round Communication rounds, T 5050 Local steps per round, E 64,128,256\64,128,256\ Per-step batch size 88 Sequence length, S 512512 Local optimizer AdamW, β1=0.9 _1=0.9, β2=0.999 _2=0.999, ϵ=10−8ε=10^-8 Local weight decay 0.010.01 Local learning rate 10−310^-3 Gradient clipping Global ℓ2 _2 norm clipping with threshold 1.01.0 Sampling Infinite RandomSampler with replacement Random seeds 0,1,2\0,1,2\ for PyTorch and NumPy The following metrics are saved for every round: • Train Perplexity (per client). During each local step, the client computes the token-level cross-entropy of the next-token objective: the logits at positions 1,…,S−11,…,S-1 are scored against the tokens at positions 2,…,S2,…,S, with S=512S=512. The pre-aggregation training perplexity of client k in round r is exp⁡(CE¯) ( CE), where CE¯ CE is the mean cross-entropy over the E local steps of the round. • Validation Perplexity. After aggregation, the global model θr+1θ^r+1 is evaluated on the server validation pool. We report exp⁡(CE¯) ( CE) with +∞+∞ saturation on overflow. Federated Algorithms. We evaluate nine federated learning algorithms, all sharing the schedule of Table 14. Throughout, pkp_k denotes the sample-count weight of client k, and since every client processes the same number of sequences per round, pk=1/Kp_k=1/K in all runs. • FedAvg [McMahan et al., 2017]. After each client completes its E local AdamW updates, the server sets θr+1=∑k=1Kpkθkr,E=:θ¯r+1θ^r+1= _k=1^Kp_k _k^r,E=: θ^r+1. This method serves as the main baseline and introduces no algorithm-specific hyperparameters. • FedProx [Li et al., 2020]. FedProx adds the uniform proximal penalty μ2‖θ−θr‖F2 μ2\|θ-θ^r\|_F^2 to the local objective, which discourages local models from drifting far from the broadcast model. Server aggregation is identical to FedAvg. • SCAFFOLD [Karimireddy et al., 2020]. SCAFFOLD corrects each client’s local update with the difference between a server control variate crc^r and the client control variate ckrc_k^r: the local gradient is shifted by λc(cr−ckr) _c\,(c^r-c_k^r) with correction scale λc _c. The control variates are refreshed every fcf_c rounds with control learning rate ηc _c, and their norm is bounded by τc _c. Aggregation of the model parameters is identical to FedAvg. • FedAdam [Reddi et al., 2021]. Clients run the unchanged AdamW local phase. The server applies Adam to the pseudo-gradient θ¯r+1−θr θ^r+1-θ^r: it maintains first- and second-moment estimates mrm^r and vrv^r with coefficients (β1s,β2s)( _1^s, _2^s) and updates θr+1=θr+αsmr/(vr+τs)θ^r+1=θ^r+ _s\,m^r/( v^r+ _s) coordinatewise, with server learning rate αs _s and offset τs _s. • FedAlign [Keçeci et al., 2025]. FedAlign aligns the Mamba2 state coordinates of the client models before averaging. In the default mode (align_mode = mamba2_ref), the B/CB/C state features of every client are matched to the current global model as the reference within each group, applying the same permutation to all coupled parameters. Optionally (align_mamba2_heads = 1), heads are matched as well, permuting every head-indexed parameter. After alignment, the client parameters are averaged with the same weights as FedAvg. • FedDyn [Acar et al., 2021]. Each client keeps a persistent dual state νkr _k^r, initialized at zero, and trains the local objective ℒk(θ)−⟨νkr,θ⟩F+α2‖θ−θr‖F2L_k(θ)- _k^r,θ _F+ α2\|θ-θ^r\|_F^2, with the two regularization terms computed in float32 outside autocast. After local training, the client sets νkr+1=νkr−α⁡(θkr,E−θr) _k^r+1= _k^r-α( _k^r,E-θ^r). The server keeps the average dual state νr+1=νr−α∑kpk(θkr,E−θr)ν^r+1=ν^r-α _kp_k( _k^r,E-θ^r) and updates θr+1=θ¯r+1−νr+1/αθ^r+1= θ^r+1-ν^r+1/α. The local optimizer remains AdamW, and the gradient of the regularized objective is clipped like every other local gradient. • FedSAM [Qu et al., 2022]. Every local step is a first-order sharpness-aware step: the client computes the gradient g at the current iterate θ, perturbs the weights by u=ℜSAMg/(‖g‖F+10−12)u= R_SAM\,g/(\|g\|_F+10^-12) with the norm taken jointly over all trainable parameters, computes the gradient at θ+uθ+u, restores θ exactly, and applies the perturbed gradient (clipped) with AdamW. Aggregation is identical to FedAvg. Because the perturbation must be formed from exact gradients across two backward passes, FedSAM runs without mixed precision at roughly twice the wall-clock cost of the other methods. • FedExP [Jhunjhunwala et al., 2023]. Local training is identical to FedAvg. With the client displacements θkr,E−θr _k^r,E-θ^r and their weighted average ∑kpk(θkr,E−θr) _kp_k ( _k^r,E-θ^r ), the server applies θr+1=θr+ηgr∑kpk(θkr,E−θr)θ^r+1=θ^r+ _g^r _kp_k ( _k^r,E-θ^r ) with the scaling factor ηgr=max⁡1,∑kpk‖Δkr‖F2/(2‖Δr‖F2+ϵg) _g^r= \1,\ _kp_k\| _k^r\|_F^2 / (2\| ^r\|_F^2+ _g ) \, a single scalar shared by all parameters, with norms taken over all floating-point tensors. The factor has no upper cap and is logged every round. We do not average the last two server iterates at evaluation time. • DiLoCo-style [Douillard et al., 2024]. Clients run the same AdamW local phase as FedAvg. The server applies an outer Nesterov step to the pseudo-gradient θ¯r+1−θr θ^r+1-θ^r: mor=βomor−1+θ¯r+1−θrm_o^r= _om_o^r-1+ θ^r+1-θ^r and θr+1=θr+αo(βomor+θ¯r+1−θr)θ^r+1=θ^r+ _o( _om_o^r+ θ^r+1-θ^r), with mo−1=0m_o^-1=0, outer learning rate αo _o, outer momentum βo _o, and per-parameter momentum buffers kept in float32. Unlike DiLoCo, whose workers keep their inner AdamW state across outer steps, the inner optimizer is re-created every round, as for all other methods. Hyperparameter Tuning. For each algorithm, we ran every value of the grid in Table 15 at every (depth,E,β,setting)(depth,E,β,setting) configuration and over three seeds. Table 15: Hyperparameter grids of the federated algorithms, run at every configuration and seed. Values shared by all members of a grid are fixed. Algorithm Hyperparameters and values FedAvg none FedProx μ∈5×10−5, 10−4, 5×10−4, 10−3, 5×10−2μ∈\5×10^-5,\,10^-4,\,5×10^-4,\,10^-3,\,5×10^-2\ SCAFFOLD fc∈4,8,16f_c∈\4,8,16\; ηc=0.01 _c=0.01; λc=0.01 _c=0.01; τc=0.1 _c=0.1 FedAdam αs∈10−3, 5×10−3, 10−2 _s∈\10^-3,\,5×10^-3,\,10^-2\; β1s=0.9 _1^s=0.9; β2s=0.99 _2^s=0.99; τs=10−3 _s=10^-3 FedAlign align_mode == mamba2_ref; align_mamba2_heads ∈0,1∈\0,1\ FedDyn α∈10−3, 3×10−3, 10−2, 3×10−2, 10−1α∈\10^-3,\,3×10^-3,\,10^-2,\,3×10^-2,\,10^-1\ FedSAM ℜSAM∈10−3, 3×10−3, 10−2, 5×10−2, 10−1 R_SAM∈\10^-3,\,3×10^-3,\,10^-2,\,5×10^-2,\,10^-1\ FedExP ϵg∈10−4, 3×10−4, 10−3, 3×10−3, 10−2 _g∈\10^-4,\,3×10^-4,\,10^-3,\,3×10^-3,\,10^-2\ DiLoCo-style αo∈0.1,0.2,0.3,0.4,0.5,0.6,0.7,1.0 _o∈\0.1,0.2,0.3,0.4,0.5,0.6,0.7,1.0\; βo=0.9 _o=0.9 The final reported metrics correspond, per configuration, to the hyperparameter value with the lowest validation perplexity averaged over the last 1010 rounds (rounds 41-50) and over the three seeds. Runs that end with a non-finite perplexity are counted as failures and are never replaced. Reporting Metrics. For every (depth,E,β,setting)(depth,E,β,setting) configuration we report: • train_ppl: the per-round, per-client training perplexity exp⁡(CE¯) ( CE) averaged over rounds and clients, then averaged over seeds (with its standard deviation). • val_ppl: the server validation perplexity exp⁡(CE¯) ( CE) averaged over all 5050 rounds, then over seeds (with its standard deviation). • val_ppl_last10: the server validation perplexity averaged over rounds 41–50, then over seeds (with its standard deviation). Ablations. We further add two ablation experiments to the main grid, both at the constant settings β=0.7β=0.7, depth 44, E=64E=64, T=50T=50 with three seeds and one fixed hyperparameter value per algorithm: FedProx μ=10−3μ=10^-3 (5×10−45×10^-4 in the backbone ablation), SCAFFOLD fc=8f_c=8, FedAdam αs=10−2 _s=10^-2, FedAlign without head matching (with head matching in the backbone ablation), FedDyn α=10−2α=10^-2, FedSAM ℜSAM=10−2 R_SAM=10^-2, FedExP ϵg=10−3 _g=10^-3, and DiLoCo-style (αo,βo)=(0.7,0.9)( _o, _o)=(0.7,0.9). The client-count ablation varies K∈4,8,16,32K∈\4,8,16,32\ on all four domain settings. Clients are assigned primary domains round-robin over the four domains of the setting, so K/4K/4 clients share each primary domain, and each client draws a fraction β of its data from its primary domain and (1−β)/3(1-β)/3 from each of the other three. Since the demand on every domain grows as K/4K/4, the per-client budget is capped at the available domain pool divided by this demand: the 7,427/1,8577,427/1,857 training/validation sequences per client at K=4K=4 become approximately 3,713/9283,713/928, 1,856/4641,856/464, and 928/232928/232 at K=8K=8, 1616, and 3232 for Settings 1 and 4, and slightly fewer for Settings 2 and 3 (which are limited by by the size of c4). The backbone ablation keeps K=4K=4 and Setting 1 and replaces some or all Mamba2 mixers of the depth-44 model by attention layers. The six backbones are described in the Models subsection. G.5 Model Merging Framework The merging framework is a two-phase pipeline. In the first phase, we fine-tune the pretrained base θ0θ^0 on each of the four downstream tasks independently, producing four task specialists θm _m. In the second phase, we merge the specialists with each configured method and evaluate the merged state on every task’s evaluation split. Per-Task Local Fine-Tuning For each task m∈ag_news,dbpedia_14,trec,imdbm∈\ ag\_news,\, dbpedia\_14,\, trec,\, imdb\ the pretrained base θ0θ^0 is fine-tuned on the task’s training pool with the causal-LM cross-entropy objective restricted to label-token positions. Specialists are checkpointed at every epoch in ”Checkpoint epochs” so that phase 2 can be repeated at multiple training horizons without re-paying the fine-tuning cost. Table 16 lists the optimization parameters. Table 16: Per-task fine-tuning recipe (identical across tasks and seeds). Component Value Optimizer AdamW Learning rate 2×10−42× 10^-4 Weight decay 0.0 Gradient clipping ℓ2 _2 norm, threshold 1.0 Train batch size 4 for mamba2-130m, mamba2-370m, 8 for mamba2-780m Eval batch size 4 LR schedule linear warmup + cosine decay Warmup ratio 0.05 (5% of total steps) Total epochs 4 Checkpoint epochs 1,2,4\1,2,4\ Compute / storage dtype bfloat16 Seeds per (model, task) 0,1,2\0,1,2\ Sequence length 256 tokens The cross-entropy loss is masked to label-token positions only, prompt and template tokens are ignored. The training-pool subsample varies with the seed (so each seed sees a different 40k subset of the underlying task), while the evaluation-pool subsample is held fixed across seeds (eval seed = 99,991) so that all reported accuracies are comparable across seeds and methods. G.6 Reproducibility Deterministic Data and Seed Handling. All three frameworks share a single PyTorch + NumPy seed. The federated harness deterministically caches the tokenized RedPajama subsets to a single binary cache file shared across all algorithms, settings, and seeds. The merging harness deterministically samples train and eval pools from each downstream dataset, with the eval seed (99,991) held fixed across all training seeds so that all reported accuracies are directly comparable. Hardware and Runtime. All reported experiments were run on single-GPU workers using NVIDIA A100-40GB and NVIDIA H100 GPUs. The A100-40GB workers were used for prototyping, debugging, and smaller validation runs, as well as the synthetic dataset generation and validation experiments. The remaining main sweeps of federated and model merging experiments were run on the H100 workers. Each run used one GPU worker, with host CPU/RAM used only for data loading, tokenization-cache access, logging, and checkpoint I/O. The tokenized RedPajama subsets were cached once and reused across algorithms, settings, and seeds. For federated language-modeling experiments, the dominant cost was local client training. Runtime scaled approximately linearly with the number of local steps E. In our setup, single federated runs ranged from approximately 3030 minutes for E=64E=64 to at most roughly 2 hours for the longest E=256E=256 runs on an H100. All federated algorithms used the same base hardware and data pipeline. Their runtime differences were small relative to the effect of E. Naturally, the 1616-layer model scaled the computation compared to the 44-layer one. For model merging, each Phase-1 task-expert fine-tuning configuration required up to roughly 8 hours on a single H100, depending on checkpoint size. Appendix H Additional Experimental Results and Ablation Studies H.1 Single Layer Synthetic Distributed Learning Note that we validate the single-layer statements only. The multi-layer bound of Theorem C.3 composes the same per-layer constants with the layerwise sensitivities κq _q of (303): the contribution of layer q is the single-layer constant of that layer (with its parameters adjusted to the multi-layer setting) multiplied by ∏j>qκj _j>q _j over the layers above it. When κj≥1 _j≥ 1, these factors do not reduce the per-layer bounds, but this does not establish greater relative looseness. Useful layers can also contract perturbations. Numerical evaluation is restricted to the single-layer theory. Validity and tightness. Table 17 reports, per ρ, the number of evaluated cells, parameter points and sequences, the violation counts for the configured and tightened bounds (all zero), the median utilization with confidence intervals, the largest utilization and the median realized ρ ρ. Table 18 gives the same statistics along training. Figure 3 shows the utilization along training, one curve per ρ, and Figure 4 the tightened-constant companion of Figure 2(a) together with the full per-sequence utilization distributions: at every ρ the distribution ends at least two orders of magnitude below the bound. 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 SSMG_SSM (conf.) or SSMtightG^tight_SSM (tight.). The row marked ∗ is the one ρ=0.3ρ=0.3 cell whose realized step δ δ differs from the other eight and is grouped separately by the evaluation. ρ cells points sequences viol. (conf./tight.) median util. (conf.) median util. (tight.) max util. (tight.) median ρ ρ 0.3∗ 1 18 18,432 0 / 0 7.6×10−57.6× 10^-5 7.2×10−47.2× 10^-4 8.3×10−38.3× 10^-3 0.0602 0.3 8 144 147,456 0 / 0 2.0×10−62.0× 10^-6 1.2×10−41.2× 10^-4 4.1×10−34.1× 10^-3 0.0345 0.5 9 162 165,888 0 / 0 1.6×10−61.6× 10^-6 3.7×10−53.7× 10^-5 5.1×10−35.1× 10^-3 0.146 0.7 9 162 165,888 0 / 0 7.7×10−77.7× 10^-7 1.6×10−51.6× 10^-5 2.1×10−32.1× 10^-3 0.369 0.9 9 162 165,888 0 / 0 2.3×10−72.3× 10^-7 4.1×10−64.1× 10^-6 1.0×10−31.0× 10^-3 0.746 0.95 9 162 165,888 0 / 0 9.1×10−89.1× 10^-8 2.0×10−62.0× 10^-6 5.7×10−45.7× 10^-4 0.868 0.99 9 162 165,888 0 / 0 3.2×10−83.2× 10^-8 7.3×10−77.3× 10^-7 1.6×10−41.6× 10^-4 0.97 0.995 9 162 165,888 0 / 0 1.2×10−81.2× 10^-8 7.6×10−77.6× 10^-7 2.4×10−42.4× 10^-4 0.982 0.999 9 162 165,888 0 / 0 5.3×10−105.3× 10^-10 6.2×10−76.2× 10^-7 1.2×10−41.2× 10^-4 0.996 all 72 1,296 1,327,104 0 / 0 1.2×10−71.2× 10^-7 3.5×10−63.5× 10^-6 8.3×10−38.3× 10^-3 0.823 Table 18: Utilization along training (single layer). checkpoint points median util. (conf.) median util. (tight.) max util. (tight.) median SSMG_SSM median gradient norm frozen_teacher 216 2.7×10−82.7× 10^-8 4.0×10−74.0× 10^-7 1.8×10−41.8× 10^-4 9.0×1069.0× 10^6 0.177 learner_final 216 1.2×10−71.2× 10^-7 4.4×10−64.4× 10^-6 3.5×10−33.5× 10^-3 6.5×1056.5× 10^5 0.0612 learner_initialization 216 3.4×10−73.4× 10^-7 1.9×10−51.9× 10^-5 8.3×10−38.3× 10^-3 1.1×1051.1× 10^5 0.0293 learner_intermediate 648 1.4×10−71.4× 10^-7 5.8×10−65.8× 10^-6 7.4×10−37.4× 10^-3 4.0×1054.0× 10^5 0.0438 (a) Configured constants (b) Tightened constants Figure 3: Median bound utilization along training, one curve per ρ (colour), with the frozen teacher at the last position. The bound is tightest at initialization and loosest at the frozen teacher at every ρ. (a) Gradient versus SSMtightG^tight_SSM (b) Per-sequence utilization Figure 4: (4(a)) Largest measured gradient norm over the test split against the tightened SSMtightG^tight_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 ρ (one in eight sequences shown), the vertical line is the bound. Anatomy of the bound and of the gradient. Table 19 lists the median share of the squared gradient norm carried by each parameter block at the final learner checkpoint and at the frozen teacher, next to the share of SSMG_SSM carried by its recurrent-Jacobian term. 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. checkpoint ρ WΔW_ WBW_B WCW_C A recurrent share of SSMG_SSM final (step 2000) 0.3 0.148 0.428 0.431 0.000 0.994 0.5 0.145 0.434 0.433 0.000 0.996 0.7 0.121 0.432 0.428 0.002 0.997 0.9 0.044 0.369 0.415 0.039 0.999 0.95 0.039 0.387 0.429 0.079 0.999 0.99 0.006 0.187 0.204 0.600 1.000 0.995 0.005 0.096 0.122 0.780 1.000 0.999 0.003 0.067 0.072 0.852 1.000 frozen teacher 0.3 0.050 0.594 0.331 0.000 0.999 0.5 0.077 0.568 0.404 0.000 0.999 0.7 0.065 0.548 0.350 0.003 0.999 0.9 0.043 0.482 0.301 0.111 1.000 0.95 0.028 0.354 0.124 0.325 1.000 0.99 0.002 0.393 0.015 0.581 1.000 0.995 0.002 0.420 0.008 0.570 1.000 0.999 0.001 0.362 0.004 0.631 1.000 Dimension. Table 20 and Figure 2(d) give the log-log exponents of the configured bound, the tightened bound and the measured gradient in D (at fixed N, seed and ρ) and in N at the final learner checkpoint. The certified exponents follow the ND2N\,D^2-type structure of KWK_W, the measured gradient is flat in N and decreasing in D. 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). ρ G in D SSMtightG^tight_SSM in D ‖∇ℒ‖\| \| in D G in N SSMtightG^tight_SSM in N ‖∇ℒ‖\| \| in N 0.3 1.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.5 1.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.7 1.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.9 1.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.95 1.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.99 1.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] 0.995 2.93 [0.71, 3.40] 2.09 [0.12, 2.94] -0.30 [-1.39, -0.21] 1.70 [0.99, 2.15] 1.55 [1.10, 2.48] 0.50 [-0.09, 0.56] 0.999 2.90 [1.17, 3.85] 2.33 [0.91, 3.71] 0.16 [-0.06, 0.45] 1.27 [0.63, 2.20] 1.34 [0.81, 2.75] 0.17 [-0.41, 0.49] Realized stability and training. Table 21 and Figure 5 demonstrate that the bounds use learners that generally outperform uniform prediction and approach the teachers’ test Negative Log Likelihood (NLL). At short memory (ρ far away from 1), the teachers’ small predictive advantage over uniform makes normalized gap recovery sensitive to small absolute errors, while the learners remain well inside the stability limit. At long memory (ρ approaches 1), learners nearly match the teachers and train close to that limit. 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. ρ runs recovered fraction median excess NLL median gap recovered [95% CI] median final ρ ρ 0.3 27 0.93 0.0180 0.80 [0.52, 0.90] 0.031 0.5 27 0.93 0.0212 0.79 [0.49, 0.82] 0.134 0.7 27 0.93 0.0203 0.81 [0.40, 0.87] 0.357 0.9 27 0.93 0.0187 0.69 [0.64, 0.96] 0.724 0.95 27 1.00 0.0179 0.74 [0.52, 0.99] 0.855 0.99 27 1.00 0.0027 0.95 [0.42, 1.03] 0.966 0.995 27 1.00 0.0004 1.00 [0.75, 1.02] 0.981 0.999 27 1.00 -0.0019 1.04 [0.95, 1.05] 0.996 (a) Gap recovered by the learners (b) Predictive gap of the teachers Figure 5: (5(a)) Fraction of the uniform-to-teacher NLL gap recovered by the learners, per ρ. (5(b)) The teachers’ predictive gap over the uniform baseline as a percentage of the uniform NLL, per cell. Federated ingredients. Figure 6 shows how the quantities that enter Theorem B.5 scale with the federated hyperparameters. Client drift, the end-of-round displacement ‖Wkr,E−Wr‖F\|W_k^r,E-W^r\|_F, grows as E0.87E^0.87 while the initialization-based reference radius R=ηESSM(W0)R=η EG_SSM(W^0) grows as E1E^1 (Figure 6(a)). The client disagreement ∑kpk‖∇ℒk(Wr)−∇f(Wr)‖F2 _kp_k\| _k(W^r)-∇ f(W^r)\|_F^2, the quantity whose supremum ΓSSM,R2 ^2_SSM,R Lemma B.4 bounds, grows as K1.48K^1.48 (Figure 6(b)) and as (1−ρ)−1.15(1-ρ)^-1.15, whereas drift is flat in ρ and ‖∇f‖F2\|∇ f\|_F^2 grows only as (1−ρ)−0.2(1-ρ)^-0.2 (Figure 6(c)), so disagreement grows faster than the squared gradient norm, consistent with the lemma’s SSM-sensitivity dependence. Figure 7 shows the per-round curves behind these tail means. Figure 8 examines the heterogeneity lemma: measured disagreement is 10−1110^-11 to 10−1410^-14 of 4SSM2ℋTV4G^2_SSMH_TV (Figure 8(a)). The sequence-level squared-TV aggregate is 0.560.56 at K=4K=4 (0.250.25 to 0.940.94 across K) and unchanged by the dominant-token partition, whereas a token-level squared-TV analogue rises from 1.6×10−51.6× 10^-5 to 2.8×10−32.8× 10^-3 and disagreement by 226×226× (Figure 8(b)). Thus, the sequence-level term is insensitive to this partition; certifying a token-level or moment-based alternative requires a separate bound. (a) Client drift against E (b) Dissimilarity against K (c) Base setting against 1−ρ1-ρ Figure 6: Federated scaling directions (medians with bootstrap 95% intervals). (6(a)) End-of-round client drift against E for FedAvg and FedProx on the core cells, with the initialization-based reference radius R=ηESSM(W0)R=η E\,G_SSM(W^0). (6(b)) Client dissimilarity ∑kpk‖∇ℒk(Wr)−∇f(Wr)‖F2 _kp_k\| _k(W^r)-∇ f(W^r)\|_F^2 against K on the core cells. (6(c)) ∑kpk‖∇ℒk(Wr)−∇f(Wr)‖F2 _kp_k\| _k(W^r)-∇ f(W^r)\|_F^2, drift and ‖∇f‖F2\|∇ f\|_F^2 against 1−ρ1-ρ at the base setting (all 72 cells). (a) E=1E=1 (b) E=4E=4 (c) E=16E=16 (d) E=64E=64 Figure 7: Per-round gradient (top) and client drift (bottom) on the core cells for each E (a) Heterogeneity lemma (b) Total-variation terms 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. What drives the bound. Table 22 and Figure 9 give signed covariance attribution across the bank. The recurrent-history term carries at least 99%99\% of the bound, so log10⁡SSM _10G_SSM is the sum of log-factors ρ/(1−ρ)2ρ/(1-ρ)^2, NDND, ‖WC‖2\|W_C\|_2, KBK_B, KΔK_ , and KΔ+D‖A‖2K_ + D\|A\|_2, plus a residual below 0.10.1 decades. Each share is the covariance of a log-factor with log-bound magnitude divided by the latter’s variance. The memory factor spans 6.26.2 decades and contributes 93%93\%. Dimension, readout and step contribute 7%7\%, 6%6\% and 4%4\%, while the transition factor contributes −13%-13\%, offsetting memory growth. This offset helps explain the exponent above −2-2 in Figure 2(b). Teacher initialization gives ‖A‖2=γ∑n(1+0.25|zn|)2\|A\|_2=γ _n(1+0.25|z_n|)^2 with γ=−lnρ/δγ=- ρ/δ. For log-utilization, memory contributes 102%102\% and measured gradient growth −18%-18\%. Recurrence therefore dominates variation in tightness across this bank, not necessarily absolute looseness. Table 22: What drives the bound across the bank. log10⁡SSM _10G_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 log10⁡SSM _10G_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 ‖∇ℒ‖F/SSM\| \|_F/G_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ρ→ 1. factor range (decades) std (decades) share of Var⁡(log10⁡SSM)Var( _10G_SSM) share of Var(log10util.)Var( _10\,util.) memory ρ/(1−ρ)2ρ/(1-ρ)^2 6.2 2.06 93% 102% dimension NDN\,D 1.2 0.35 7% 12% readout ‖WC‖2\|W_C\|_2 1.5 0.31 6% 6% input KBK_B 1.8 0.28 3% 3% step KΔK_ 0.8 0.18 4% 6% transition KΔ+D‖A‖2K_ + D\,\|A\|_2 2.0 0.50 -13% -11% remaining terms of SSMG_SSM 0.1 0.01 0% 0% measured ‖∇ℒ‖F\| \|_F 2.7 0.62 – -18% (a) What drives SSMG_SSM (b) What drives the utilization ‖∇ℒ‖F/SSM\| \|_F/G_SSM Figure 9: Additive attribution of the variance of (9(a)) log10⁡SSM _10G_SSM and (9(b)) log10 _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. H.2 Practical Distributed Learning In this subsection we provide supplemental results to the Section 5 Mamba2 federated learning experiments. Every table reports the three metrics defined in Appendix G: the client training perplexity, the server validation perplexity (both averaged over all 5050 rounds), and the server validation perplexity averaged over the last 1010 rounds. Appendix H.2.1 evaluates the 99 federated algorithms in all the 7272 main configurations in details, demonstrating how the different local rounds E, model depths, heterogeneity degrees β, and data compositions affect their performance. The reported standard deviation measures the per-seed average across configurations for 3 seeds (not the perplexity spread across configurations). Appendices H.2.2 and H.2.3 provide two ablation studies, on different client counts, and different model backbones respectively. Both use fixed hyperparameters listed in the Ablations paragraph of Appendix G. H.2.1 Detailed Results Across the 72 Main Configurations Tables 23-26 report the nine algorithms averaged over the 7272 configurations along each factor of the grid. Table 23: Effect of heterogeneity. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms on Mamba2 at every β, averaged over the 2424 configurations. β=0β=0 β=0.7β=0.7 β=1β=1 Algorithm Train Val Last-10 Train Val Last-10 Train Val Last-10 FedAvg 75.4 ± 0.5 63.7 ± 0.3 45.4 ± 0.2 99.3 ± 0.2 66.7 ± 0.6 45.8 ± 0.3 143.5 ± 0.7 136.5 ± 3.8 60.7 ± 0.5 FedProx 69.0 ± 0.2 57.6 ± 0.3 36.9 ± 0.4 88.0 ± 0.4 60.4 ± 0.4 36.5 ± 0.2 119.3 ± 0.4 121.3 ± 2.9 40.2 ± 0.3 SCAFFOLD 78.3 ± 0.5 66.3 ± 0.2 46.9 ± 0.1 102.4 ± 0.4 67.3 ± 0.3 45.2 ± 0.1 149.0 ± 0.5 131.6 ± 3.6 49.7 ± 0.2 FedAdam 88.4 ± 0.4 164.8 ± 1.6 53.0 ± 0.3 112.7 ± 0.4 160.0 ± 2.1 50.8 ± 0.1 157.2 ± 0.6 355.2 ± 13.8 61.7 ± 0.3 FedAlign 75.3 ± 0.4 63.6 ± 0.2 45.2 ± 0.0 99.2 ± 0.4 66.6 ± 0.3 45.7 ± 0.2 143.5 ± 0.5 135.3 ± 5.4 60.1 ± 0.4 FedDyn 78.9 ± 0.1 106.4 ± 2.1 44.1 ± 0.5 95.5 ± 0.2 120.1 ± 3.9 41.6 ± 0.1 127.1 ± 0.0 814.4 ± 50.2 40.0 ± 0.2 FedSAM 75.4 ± 0.4 63.7 ± 0.3 45.2 ± 0.1 99.3 ± 0.4 66.7 ± 0.3 45.7 ± 0.1 144.5 ± 0.7 135.5 ± 3.4 59.8 ± 0.3 FedExP 75.3 ± 0.4 63.6 ± 0.2 45.2 ± 0.1 97.6 ± 0.4 66.5 ± 0.4 45.2 ± 0.2 137.8 ± 0.6 154.9 ± 7.1 68.9 ± 0.5 DiLoCo-style 72.4 ± 0.6 66.4 ± 0.4 46.1 ± 0.3 96.2 ± 0.5 71.7 ± 0.4 46.2 ± 0.2 140.5 ± 0.6 202.2 ± 9.8 61.3 ± 0.2 Table 24: Effect of local steps. Training, validation, and last-10-round validation perplexity of the nine federated learning algorithms on Mamba2 at every E, averaged over the 2424 configurations. E=64E=64 E=128E=128 E=256E=256 Algorithm Train Val Last-10 Train Val Last-10 Train Val Last-10 FedAvg 151.2 ± 0.5 108.5 ± 1.7 63.9 ± 0.2 100.9 ± 0.5 85.5 ± 1.1 48.7 ± 0.2 66.1 ± 0.3 72.9 ± 2.0 39.2 ± 0.4 FedProx 137.7 ± 0.3 106.6 ± 1.7 48.6 ± 0.2 85.1 ± 0.4 74.9 ± 0.9 35.6 ± 0.4 53.6 ± 0.3 57.8 ± 1.0 29.4 ± 0.2 SCAFFOLD 150.8 ± 0.4 105.6 ± 1.7 61.0 ± 0.3 104.5 ± 0.3 85.3 ± 0.7 45.4 ± 0.3 74.5 ± 0.4 74.3 ± 2.0 35.4 ± 0.0 FedAdam 167.3 ± 0.5 261.5 ± 5.6 69.6 ± 0.5 113.0 ± 0.4 217.6 ± 6.6 53.1 ± 0.1 78.1 ± 0.3 200.9 ± 5.2 42.8 ± 0.1 FedAlign 151.2 ± 0.5 108.3 ± 2.1 63.6 ± 0.2 100.8 ± 0.4 85.1 ± 1.7 48.4 ± 0.3 66.0 ± 0.3 72.0 ± 2.3 39.0 ± 0.2 FedDyn 149.6 ± 0.1 257.9 ± 5.0 56.2 ± 0.3 90.6 ± 0.2 351.7 ± 43.0 38.8 ± 0.4 61.3 ± 0.1 431.3 ± 48.0 30.7 ± 0.0 FedSAM 151.3 ± 0.5 109.0 ± 1.5 63.8 ± 0.3 101.1 ± 0.4 85.6 ± 0.8 48.3 ± 0.1 66.8 ± 0.5 71.3 ± 1.8 38.6 ± 0.2 FedExP 147.1 ± 0.6 111.7 ± 1.2 65.3 ± 0.2 98.6 ± 0.5 91.5 ± 1.1 51.9 ± 0.3 65.0 ± 0.4 81.8 ± 5.6 42.1 ± 0.4 DiLoCo-style 144.6 ± 0.5 140.8 ± 5.4 63.5 ± 0.3 97.5 ± 0.6 113.4 ± 4.7 49.3 ± 0.3 67.0 ± 0.5 86.1 ± 3.2 40.9 ± 0.1 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. Setting 1 Setting 2 Setting 3 Setting 4 Algorithm Train Val Last-10 Train Val Last-10 Train Val Last-10 Train Val Last-10 FedAvg 45.1 ± 0.2 53.2 ± 1.5 25.2 ± 0.3 107.3 ± 0.0 96.5 ± 4.2 51.7 ± 0.6 140.6 ± 1.0 91.4 ± 3.3 56.9 ± 0.1 131.3 ± 0.5 114.8 ± 1.1 68.7 ± 0.8 FedProx 42.3 ± 0.2 46.4 ± 0.2 19.0 ± 0.2 93.8 ± 0.3 86.4 ± 0.3 38.9 ± 0.3 117.7 ± 0.6 87.3 ± 3.1 43.0 ± 0.3 114.6 ± 0.5 98.9 ± 1.7 50.5 ± 0.9 SCAFFOLD 45.9 ± 0.2 51.6 ± 1.5 22.6 ± 0.4 112.0 ± 0.0 95.6 ± 3.4 47.5 ± 0.2 146.5 ± 0.7 93.1 ± 3.2 55.6 ± 0.1 135.2 ± 0.5 113.3 ± 1.3 63.3 ± 0.7 FedAdam 51.4 ± 0.1 169.0 ± 4.3 27.2 ± 0.2 118.2 ± 0.4 260.3 ± 6.6 54.2 ± 0.2 158.0 ± 1.1 237.0 ± 6.4 62.2 ± 0.2 150.2 ± 0.2 240.3 ± 7.6 77.1 ± 1.3 FedAlign 45.0 ± 0.0 50.8 ± 1.6 24.8 ± 0.2 107.2 ± 0.3 96.7 ± 3.4 51.4 ± 0.3 140.6 ± 0.9 91.8 ± 2.7 56.7 ± 0.1 131.3 ± 0.6 114.6 ± 1.2 68.5 ± 0.9 FedDyn 46.1 ± 0.0 339.0 ± 57.0 20.8 ± 0.2 101.3 ± 0.2 502.0 ± 35.6 42.3 ± 0.3 127.1 ± 0.5 248.9 ± 70.4 47.7 ± 0.2 127.4 ± 0.2 297.9 ± 23.3 56.7 ± 0.7 FedSAM 45.4 ± 0.2 51.8 ± 1.6 25.0 ± 0.2 107.8 ± 0.3 96.2 ± 3.6 50.9 ± 0.2 141.0 ± 0.9 91.9 ± 2.2 56.6 ± 0.1 131.6 ± 0.5 114.8 ± 1.2 68.6 ± 0.9 FedExP 43.9 ± 0.3 60.7 ± 3.9 27.0 ± 0.1 104.8 ± 0.2 105.5 ± 3.8 55.0 ± 0.3 137.3 ± 1.0 97.5 ± 4.2 60.2 ± 0.1 128.3 ± 0.5 116.3 ± 1.1 70.2 ± 0.7 DiLoCo-style 43.2 ± 0.3 62.8 ± 2.5 25.6 ± 0.3 104.6 ± 0.2 130.7 ± 9.2 52.3 ± 0.3 137.6 ± 1.1 110.2 ± 3.2 58.4 ± 0.5 126.9 ± 0.5 149.9 ± 5.4 68.6 ± 0.8 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. Depth 44 (14.714.7M) Depth 1616 (20.320.3M) Algorithm Train Val Last-10 Train Val Last-10 FedAvg 109.8 ± 0.5 88.8 ± 0.5 54.8 ± 0.3 102.3 ± 0.4 89.1 ± 2.6 46.4 ± 0.2 FedProx 93.2 ± 0.4 79.0 ± 1.3 39.5 ± 0.4 91.0 ± 0.3 80.5 ± 1.2 36.2 ± 0.1 SCAFFOLD 112.8 ± 0.4 88.0 ± 1.4 50.1 ± 0.0 107.0 ± 0.4 88.8 ± 2.0 44.5 ± 0.2 FedAdam 122.6 ± 0.3 235.1 ± 4.1 59.8 ± 0.3 116.2 ± 0.5 218.2 ± 7.4 50.5 ± 0.3 FedAlign 109.7 ± 0.5 89.1 ± 1.9 54.4 ± 0.2 102.4 ± 0.4 87.9 ± 2.2 46.3 ± 0.3 FedDyn 100.2 ± 0.2 398.3 ± 43.9 42.8 ± 0.3 100.8 ± 0.3 295.7 ± 15.5 41.0 ± 0.1 FedSAM 110.1 ± 0.6 89.8 ± 1.7 54.3 ± 0.1 102.7 ± 0.4 87.5 ± 1.1 46.2 ± 0.4 FedExP 107.5 ± 0.4 93.1 ± 2.5 57.6 ± 0.4 99.6 ± 0.6 96.9 ± 2.6 48.5 ± 0.3 DiLoCo-style 108.8 ± 0.6 112.2 ± 0.7 56.3 ± 0.1 97.3 ± 0.5 114.7 ± 6.8 46.2 ± 0.4 Comments. The results averages show a consistent overall ranking. FedProx achieves the lowest mean perplexity in 3434 of the 3636 metric columns. The two exceptions are last-10 validation perplexity at β=1β=1, where FedDyn is marginally lower (40.040.0 versus 40.240.2), and all-round validation perplexity at E=64E=64, where SCAFFOLD is lower (105.6105.6 versus 106.6106.6). On the last-10 metric, FedProx, FedDyn, and SCAFFOLD rank first, second, and third, respectively, at every local-step count, domain setting, and model depth in Tables 24-26. Heterogeneity produces the clearest changes in relative performance. FedProx reduces last-10 perplexity relative to FedAvg by approximately 19%19\%, 20%20\%, and 34%34\% at β=0β=0, 0.70.7, and 11, respectively. SCAFFOLD moves from eighth place at β=0β=0 (46.946.9 versus 45.445.4 for FedAvg) to third at β=1β=1, where it improves on FedAvg by 18%18\%. FedDyn narrowly overtakes FedProx at β=1β=1, whereas FedExP moves from near parity with FedAvg at β≤0.7β≤ 0.7 to last place at β=1β=1, with 13.5%13.5\% higher last-10 perplexity. Across the same increase in β, last-10 perplexity rises by approximately 3232-34%34\% for FedAvg, FedAlign, FedSAM, and DiLoCo-style, by 52%52\% for FedExP, and by 16%16\% for FedAdam. The corresponding increases are only 9%9\% for FedProx and 6%6\% for SCAFFOLD. FedDyn instead improves by approximately 9%9\%, although this comparison uses a separately selected α at each β. Increasing the number of local steps strengthens the relative advantages of FedDyn and SCAFFOLD. Their reductions in last-10 perplexity relative to FedAvg grow from approximately 12%12\% and 4.5%4.5\% at E=64E=64 to 22%22\% and 10%10\% at E=256E=256, respectively. FedProx maintains a comparatively uniform reduction of 2424-27%27\% across the three values of E. These trends are qualitatively consistent with the role of the E⁡(E−1)E(E-1) drift term in the bounds and the benefits of radius control or heterogeneity correction discussed in our analysis. Naturally, larger E results in longer training, so all nine algorithms achieve approximately 3535-45%45\% lower last-10 perplexity when E increases from 6464 to 256256 at fixed T=50T=50 communication rounds. The remaining methods generally remain close to, or underperform, FedAvg on the last-10 metric. FedAlign and FedSAM stay within 2%2\% of FedAvg throughout the factor tables, while DiLoCo-style stays within 5%5\%. FedAdam has approximately 9%9\% higher last-10 perplexity at every local-step count and model depth, with gaps of approximately 55-12%12\% across domain settings. All-round validation perplexity provides an important contrast to these last-10 results. Despite its strong late-training performance, FedDyn has substantially higher all-round validation perplexity than FedAvg, as does FedAdam. Across the factor tables, their all-round averages are approximately 1.71.7-6.46.4 times those of FedAvg, reaching 814.4814.4 for FedDyn and 355.2355.2 for FedAdam at β=1β=1. The round-wise results place these averages in context: after round 11, FedDyn and FedAdam have validation perplexities approximately 88 and 99 times that of FedAvg, respectively, using geometric means across configurations. FedDyn subsequently falls below FedAvg from round 44 onward and finishes 18%18\% lower, whereas FedAdam remains above FedAvg throughout training and finishes 9%9\% higher. Domain composition and model depth change the absolute perplexity levels without changing the ordering of the three leading methods. Setting 1 yields the lowest last-10 perplexity and Setting 4 the highest for every algorithm. For FedAvg, these values differ by approximately a factor of 2.72.7. FedProx improves on FedAvg by approximately 2424-26%26\% in each setting. Increasing depth from 44 to 1616 reduces last-10 perplexity by approximately 1515-18%18\% for FedAvg, FedAlign, FedSAM, FedExP, DiLoCo-style, and FedAdam, compared with 11%11\% for SCAFFOLD, 8%8\% for FedProx, and 4%4\% for FedDyn. Consequently, the relative improvements over FedAvg narrow from 28%28\% to 22%22\% for FedProx, from 22%22\% to 12%12\% for FedDyn, and from 9%9\% to 4%4\% for SCAFFOLD. H.2.2 Client-Count Ablation Table 27 reports the client-count ablation described in Appendix G. We consider K∈4,8,16,32K∈\4,8,16,32\ client counts at fixed β=0.7β=0.7, depth 44, and E=64E=64, averaged over the four domain settings, with K/4K/4 clients sharing each primary domain and every client holding a 4/K4/K share of the same domain pools. Table 27: Client-count ablation. Training, validation, and last-10-round validation perplexity of the nine algorithms with K∈4,8,16,32K∈\4,8,16,32\ clients at β=0.7β=0.7, depth 44, E=64E=64, averaged over the four domain settings. K=4K=4 K=8K=8 K=16K=16 K=32K=32 Algorithm Train Val Last-10 Train Val Last-10 Train Val Last-10 Train Val Last-10 FedAvg 144.4 ± 0.3 89.7 ± 0.9 64.2 ± 0.4 141.4 ± 0.4 87.3 ± 0.5 61.6 ± 0.2 136.9 ± 0.6 86.3 ± 0.7 60.9 ± 0.3 129.0 ± 0.3 86.2 ± 0.6 61.1 ± 0.2 FedProx 131.6 ± 0.2 82.4 ± 0.8 49.0 ± 0.2 129.1 ± 0.5 80.6 ± 0.3 47.3 ± 0.3 125.2 ± 0.4 79.9 ± 0.5 46.8 ± 0.2 118.5 ± 0.3 80.5 ± 0.4 47.6 ± 0.1 SCAFFOLD 143.7 ± 0.2 87.6 ± 0.9 62.4 ± 0.4 139.5 ± 0.5 83.8 ± 0.5 58.7 ± 0.2 134.2 ± 0.7 81.2 ± 0.6 55.7 ± 0.4 126.8 ± 0.4 79.7 ± 0.5 53.0 ± 0.2 FedAdam 159.9 ± 0.1 187.1 ± 2.0 69.5 ± 0.6 156.2 ± 0.6 183.0 ± 1.2 67.0 ± 0.4 151.3 ± 0.5 181.7 ± 2.0 66.1 ± 0.4 143.7 ± 0.2 182.8 ± 1.7 66.5 ± 0.5 FedAlign 144.4 ± 0.2 89.7 ± 0.9 64.2 ± 0.4 141.4 ± 0.4 87.3 ± 0.5 61.6 ± 0.2 136.9 ± 0.6 86.3 ± 0.7 60.9 ± 0.3 129.0 ± 0.3 86.2 ± 0.6 61.1 ± 0.2 FedDyn 164.5 ± 0.4 161.8 ± 15.9 59.4 ± 0.5 163.0 ± 0.7 156.9 ± 5.7 54.1 ± 0.4 162.2 ± 0.2 151.4 ± 10.3 51.6 ± 0.4 162.0 ± 0.2 151.9 ± 11.1 50.4 ± 0.3 FedSAM 145.1 ± 0.4 90.3 ± 0.8 64.3 ± 0.2 142.2 ± 0.5 88.1 ± 0.5 62.1 ± 0.2 137.5 ± 0.7 86.9 ± 0.7 61.2 ± 0.3 129.8 ± 0.4 86.8 ± 0.6 61.5 ± 0.2 FedExP 141.7 ± 0.1 89.9 ± 1.0 63.9 ± 0.7 137.1 ± 0.5 85.9 ± 0.4 59.5 ± 0.5 131.6 ± 0.8 84.0 ± 0.8 57.7 ± 0.4 123.6 ± 0.5 83.2 ± 0.5 57.3 ± 0.4 DiLoCo-style 138.7 ± 0.2 94.6 ± 1.4 65.8 ± 0.1 134.0 ± 0.8 90.2 ± 0.8 61.0 ± 0.3 128.3 ± 0.3 87.4 ± 0.9 58.0 ± 0.6 120.8 ± 0.6 86.5 ± 0.5 56.5 ± 0.1 Comments. With per-client work held fixed, FedAvg’s last-10 validation perplexity decreases by approximately 4%4\% from K=4K=4 to K=8K=8 (64.264.2 to 61.661.6), with little further change for the larger client counts (60.960.9 and 61.161.1). The training perplexity continues to decrease, from 144144 at K=4K=4 to 129129 at K=32K=32, potentially reflecting more frequent reuse of smaller local datasets. The three leading methods in the main grid -FedProx, FedDyn, and SCAFFOLD - remain the strongest on last-10 validation perplexity at every K. FedAlign and FedSAM remain within 1%1\% of FedAvg, whereas FedAdam has 88-9%9\% higher perplexity. FedProx maintains a reduction of 2222-24%24\% compared to FedAvg across the evaluated values of K. In contrast, the reductions achieved by SCAFFOLD and FedDyn increase from 3%3\% to 13%13\% and from 7%7\% to 18%18\%, respectively, between K=4K=4 and K=32K=32. DiLoCo-style moves from 3%3\% higher perplexity than FedAvg to 8%8\% lower, while FedExP moves from approximate parity to a 6%6\% reduction. As such, although the leading methods remain unchanged, their relative gains and the ordering of the remaining methods vary with client count. One possible explanation for these method-dependent trends is that having multiple clients per primary domain at larger K produces a more stable aggregate update than the one-client-per-domain partition at K=4K=4. However, this interpretation is not directly tested by the ablation. Increasing K also changes the total client work, local dataset sizes, and data partitioning. These results therefore describe performance under the evaluated client-count protocol at a single (β,depth,E)(β,depth,E) setting, rather than isolating the effect of client count or establishing a particular algorithmic mechanism. H.2.3 Model Backbone Ablation Table 28 reports the different model backbones ablation described in Appendix G: the depth-44 model with some or all Mamba2 mixers replaced by attention layers (Table 10), trained on Setting 1 at β=0.7β=0.7, E=64E=64, and K=4K=4. 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β=0.7, E=64E=64, K=4K=4, the second header row gives the fraction of Mamba2 layers). mamba2 hybrid75 hybrid50 hybrid50alt hybrid25 attn0 SSM 1.00 SSM 0.75 SSM 0.50 SSM 0.50 SSM 0.25 SSM 0.00 Algorithm Val Last-10 Val Last-10 Val Last-10 Val Last-10 Val Last-10 Val Last-10 FedAvg 43.4 ± 0.4 30.7 ± 0.3 41.5 ± 0.8 26.5 ± 0.5 43.3 ± 0.6 29.3 ± 0.4 52.6 ± 1.2 31.3 ± 0.2 54.4 ± 1.2 32.7 ± 0.7 69.6 ± 1.1 42.3 ± 1.1 FedProx 41.9 ± 0.4 25.4 ± 0.2 39.8 ± 0.4 21.0 ± 0.2 41.4 ± 0.3 23.6 ± 0.2 50.4 ± 0.7 24.7 ± 0.4 52.2 ± 1.3 26.8 ± 0.9 67.7 ± 1.0 34.8 ± 0.9 SCAFFOLD 42.2 ± 0.4 29.4 ± 0.3 39.2 ± 1.0 22.5 ± 0.5 41.3 ± 0.7 25.4 ± 0.8 50.3 ± 1.2 27.1 ± 0.3 52.5 ± 1.2 29.6 ± 1.3 68.3 ± 1.0 41.8 ± 0.9 FedAdam 116.0 ± 1.1 33.2 ± 0.2 116.5 ± 3.4 30.7 ± 0.5 120.4 ± 2.8 32.5 ± 0.5 141.7 ± 3.9 38.8 ± 0.2 144.4 ± 2.3 40.3 ± 0.6 180.8 ± 4.6 59.4 ± 1.1 FedAlign 43.4 ± 0.4 30.7 ± 0.3 41.7 ± 1.0 26.7 ± 0.6 43.3 ± 0.6 29.3 ± 0.4 52.5 ± 1.2 31.2 ± 0.1 54.3 ± 1.1 32.7 ± 0.7 69.6 ± 1.2 42.0 ± 0.9 FedDyn 76.2 ± 6.4 29.6 ± 0.5 68.4 ± 2.1 28.2 ± 0.4 69.5 ± 5.5 27.5 ± 1.0 114.0 ± 18.4 34.8 ± 1.0 101.8 ± 16.0 35.5 ± 0.6 132.3 ± 22.7 41.6 ± 1.4 FedSAM 43.8 ± 0.4 30.8 ± 0.3 41.8 ± 0.8 26.7 ± 0.6 43.6 ± 0.7 29.6 ± 0.5 53.7 ± 0.9 32.1 ± 0.1 55.2 ± 1.4 33.4 ± 1.2 70.5 ± 0.9 42.7 ± 0.4 FedExP 43.6 ± 0.2 30.5 ± 0.0 41.0 ± 0.7 25.7 ± 0.4 43.1 ± 0.5 28.3 ± 0.3 51.8 ± 1.2 30.0 ± 0.2 53.5 ± 1.1 30.9 ± 0.7 68.9 ± 0.9 40.3 ± 1.3 DiLoCo-style 44.8 ± 0.2 31.1 ± 0.4 42.0 ± 0.9 26.3 ± 0.5 43.7 ± 0.4 29.0 ± 0.4 56.6 ± 2.6 33.5 ± 0.9 57.9 ± 0.9 35.2 ± 0.5 78.6 ± 2.5 48.1 ± 1.2 Comments. Under FedAvg the backbones order as hybrid75 (26.526.5) << hybrid50 (29.329.3) << mamba2 (30.730.7) << hybrid50alt (31.331.3) << hybrid25 (32.732.7) << attn0 (42.342.3) on the last-10 validation perplexity. Notably, one attention layer in the middle of the stack improves on the pure Mamba2 model by 14%14\%. Attention layers placed first (hybrid50alt, hybrid50) or a mostly-attention stack lose that gain, and the attention-only control is 38%38\% worse than pure Mamba2 at this model size and sequence length. The ranking of the algorithms transfers from the pure-SSM backbone to every hybrid, with two exceptions. FedProx is the best method on all six backbones (−17-17 to −21%-21\% relative to FedAvg), and FedAlign and FedSAM stay within 3%3\% of FedAvg everywhere. SCAFFOLD, in contrast, gains 1313–15%15\% on the hybrids with two or three Mamba2 layers and 10%10\% on hybrid25, but only 4%4\% on mamba2 and 1%1\% on attn0, and FedDyn at its fixed α is unstable across architectures (−2-2 to −6%-6\% on mamba2, hybrid50, and attn0, +6+6 to +11%+11\% on the other three hybrids). The server-step methods are the second exception: their penalty grows as Mamba2 layers are replaced, from +8%+8\% (FedAdam) and +1%+1\% (DiLoCo-style) on mamba2 to +40%+40\% and +14%+14\% on attn0, whereas FedExP, neutral on mamba2, is 33-6%6\% better than FedAvg on every backbone that contains attention. H.3 Model Merging This subsection reports the one-shot merging experiments referenced in Appendix F. Following the framework of Appendix G.5, the pretrained Mamba2-130M, 370M, and 780M checkpoints (Table 11) are fine-tuned separately on AG News, DBpedia, TREC, and IMDB (Table 6) with the recipe of Table 16, and the four specialists θm=θ0+m _m=θ^0+T_m saved after 11, 22, and 44 epochs are merged by weight averaging (pm=1/4p_m=1/4), Task Arithmetic (θ0+λ∑mθ^0+λ _mT_m), and TIES-Merging (trim each task vector to the fraction d of its largest-magnitude entries, elect signs, average the agreeing entries, and scale the result by λ). Every merged model is evaluated on all four tasks by causal-LM label scoring, and we report the accuracy averaged over the four tasks (mean ± standard deviation over three fine-tuning seeds), as well as the accuracies of the pretrained model and the specialists evaluated on their own and on the other tasks. Task Arithmetic and TIES-Merging are shown at the best of the evaluated λ and (λ,d)(λ,d) for each checkpoint and epoch. Tables 31 and 32 give the full sweeps, Table 30 the per-task accuracies after four epochs, and Table 33 a variant of the 130M experiment in which the specialists are fine-tuned with the FedProx proximal term. 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. Checkpoint Model 11 epoch 22 epochs 44 epochs Selected λ / (λ,d)(λ,d) Mamba2-130M Zero-shot (pretrained) 25.7 ± 0.0 – Specialists (own task) 81.8 ± 1.8 88.4 ± 0.5 89.2 ± 0.2 – Specialists (other tasks) 23.0 ± 0.2 23.6 ± 0.1 23.4 ± 0.3 – Weight averaging 34.4 ± 2.1 36.3 ± 1.7 37.4 ± 0.6 – Task Arithmetic 53.5 ± 4.6 61.9 ± 5.1 61.3 ± 3.8 11 / 11 / 0.70.7 TIES-Merging 47.0 ± 2.0 55.0 ± 3.9 56.6 ± 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) Mamba2-370M Zero-shot (pretrained) 24.2 ± 0.1 – Specialists (own task) 86.6 ± 2.3 90.7 ± 0.1 90.8 ± 0.1 – Specialists (other tasks) 24.8 ± 0.5 24.8 ± 0.3 25.4 ± 0.1 – Weight averaging 45.6 ± 1.0 49.6 ± 0.5 50.4 ± 0.4 – Task Arithmetic 61.1 ± 2.5 71.0 ± 1.7 71.8 ± 1.2 0.70.7 / 0.70.7 / 0.70.7 TIES-Merging 62.6 ± 1.6 69.8 ± 2.0 70.9 ± 1.8 (1.3,0.5)(1.3,0.5) / (1.3,0.1)(1.3,0.1) / (1.3,0.1)(1.3,0.1) Mamba2-780M Zero-shot (pretrained) 24.7 ± 0.0 – Specialists (own task) 82.4 ± 3.3 89.4 ± 0.5 90.8 ± 0.3 – Specialists (other tasks) 24.5 ± 1.0 24.6 ± 0.1 24.9 ± 0.2 – Weight averaging 39.4 ± 2.8 45.4 ± 1.4 51.7 ± 1.5 – Task Arithmetic 59.9 ± 2.8 61.9 ± 5.6 71.4 ± 2.5 0.70.7 / 0.70.7 / 0.70.7 TIES-Merging 60.3 ± 3.0 64.5 ± 7.6 72.1 ± 3.5 (1.3,0.1)(1.3,0.1) / (1.3,0.1)(1.3,0.1) / (1.3,0.1)(1.3,0.1) 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. Mamba2-130M Mamba2-370M Mamba2-780M Model AG News DBpedia TREC IMDB AG News DBpedia TREC IMDB AG News DBpedia TREC IMDB Zero-shot (pretrained) 25.0 ± 0.0 9.9 ± 0.0 16.8 ± 0.0 51.1 ± 0.2 24.3 ± 0.0 8.1 ± 0.0 13.4 ± 0.2 51.1 ± 0.1 25.3 ± 0.0 8.7 ± 0.0 15.1 ± 0.1 49.9 ± 0.1 Specialists (own task) 89.5 ± 0.3 96.4 ± 0.2 86.7 ± 0.6 83.9 ± 0.4 90.8 ± 0.2 97.6 ± 0.2 88.7 ± 0.5 86.0 ± 0.1 90.8 ± 0.1 97.9 ± 0.2 88.4 ± 1.0 86.1 ± 0.5 Weight averaging 37.9 ± 2.6 21.6 ± 3.9 34.9 ± 5.3 55.3 ± 0.7 42.4 ± 2.3 36.8 ± 0.9 63.5 ± 4.4 59.1 ± 5.4 44.4 ± 0.6 41.1 ± 3.3 54.5 ± 1.6 66.8 ± 3.2 Task Arithmetic 64.7 ± 8.5 53.3 ± 13.8 61.7 ± 7.3 65.5 ± 2.6 75.3 ± 7.0 74.4 ± 2.8 65.1 ± 2.8 72.5 ± 2.2 80.6 ± 2.3 65.9 ± 7.9 57.9 ± 7.7 81.1 ± 1.7 TIES-Merging 72.2 ± 8.4 51.7 ± 17.9 46.2 ± 13.1 56.2 ± 0.1 77.3 ± 6.8 73.4 ± 3.6 62.9 ± 3.0 69.9 ± 0.8 83.4 ± 1.4 69.5 ± 8.2 52.4 ± 11.6 83.1 ± 1.1 Table 31: Task Arithmetic scale sweep: task-averaged accuracy (%). θTA=θ0+λ∑m _TA=θ^0+λ _mT_m over the four specialists, so λ=1/M=0.25λ=1/M=0.25 would be weight averaging and λ=1λ=1 adds every task vector at full length. The 780M sweep evaluated only λ∈0.7,1.0λ∈\0.7,1.0\. Checkpoint Epochs λ=0.2λ=0.2 λ=0.3λ=0.3 λ=0.5λ=0.5 λ=0.7λ=0.7 λ=1λ=1 λ=1.3λ=1.3 Mamba2-130M 11 31.2 ± 1.2 38.1 ± 3.2 48.3 ± 3.3 52.4 ± 4.2 53.5 ± 4.6 51.7 ± 3.8 22 32.6 ± 1.1 41.3 ± 1.9 55.5 ± 2.0 61.9 ± 4.1 61.9 ± 5.1 55.8 ± 5.4 44 32.8 ± 0.8 42.7 ± 0.8 55.8 ± 2.1 61.3 ± 3.8 57.2 ± 6.2 49.3 ± 7.0 Mamba2-370M 11 40.2 ± 1.9 49.5 ± 0.6 59.4 ± 2.9 61.1 ± 2.5 58.6 ± 3.4 54.1 ± 5.6 22 43.5 ± 1.5 54.7 ± 0.4 67.6 ± 1.0 71.0 ± 1.7 68.2 ± 2.2 59.6 ± 2.7 44 44.5 ± 1.5 55.4 ± 0.5 68.6 ± 0.5 71.8 ± 1.2 66.3 ± 2.5 54.0 ± 3.3 Mamba2-780M 11 – – – 59.9 ± 2.8 58.1 ± 3.9 – 22 – – – 61.9 ± 5.6 57.9 ± 5.2 – 44 – – – 71.4 ± 2.5 69.5 ± 4.7 – Table 32: TIES-Merging sweep after 4 specialist epochs: task-averaged accuracy (%). Rows are the output scale λ applied to the task vector, columns the fraction d of largest-magnitude entries kept per task vector. The 780M sweep evaluated only λ=1.3λ=1.3 with d∈0.1,0.5d∈\0.1,0.5\ in addition to the default (λ,d)=(1.0,0.2)(λ,d)=(1.0,0.2). Checkpoint λ d=0.1d=0.1 d=0.2d=0.2 d=0.3d=0.3 d=0.5d=0.5 Mamba2-130M 0.50.5 34.7 ± 1.5 33.7 ± 0.9 32.6 ± 0.8 32.1 ± 0.8 0.80.8 46.1 ± 1.3 44.0 ± 1.1 42.2 ± 1.1 41.2 ± 1.4 11 51.7 ± 1.6 49.6 ± 1.7 47.9 ± 1.9 46.9 ± 1.3 1.31.3 56.6 ± 3.6 54.7 ± 2.6 53.7 ± 2.3 52.9 ± 1.8 Mamba2-370M 0.50.5 45.4 ± 0.6 46.4 ± 0.3 45.0 ± 1.5 42.0 ± 1.3 0.80.8 63.0 ± 1.7 62.6 ± 0.8 60.9 ± 1.1 57.5 ± 1.3 11 68.0 ± 1.7 67.8 ± 1.2 67.1 ± 1.4 64.7 ± 1.8 1.31.3 70.9 ± 1.8 70.4 ± 1.7 70.6 ± 1.5 70.2 ± 1.6 Mamba2-780M 11 – 67.8 ± 4.2 – – 1.31.3 72.1 ± 3.5 – – 67.2 ± 3.3 Table 33: Weight averaging of proximally fine-tuned specialists (Mamba2-130M): task-averaged accuracy (%). Each specialist is fine-tuned with the FedProx local objective ℒm(θ)+μ2‖θ−θ0‖F2L_m(θ)+ μ2\|θ-θ^0\|_F^2 for 11, 22, or 44 epochs and the four specialists are merged by weight averaging Specialists (own task) Weight averaging μ 11 epoch 22 epochs 44 epochs 11 epoch 22 epochs 44 epochs 00 (reference) 81.8 ± 1.8 88.4 ± 0.5 89.2 ± 0.2 34.4 ± 2.1 36.3 ± 1.7 37.4 ± 0.6 10−510^-5 82.6 ± 3.0 88.1 ± 0.9 88.9 ± 0.6 33.3 ± 0.9 35.7 ± 0.6 37.3 ± 0.6 10−410^-4 82.5 ± 3.5 87.6 ± 0.3 88.1 ± 0.3 32.5 ± 2.0 33.1 ± 1.1 34.0 ± 0.8 10−310^-3 80.6 ± 2.6 85.3 ± 0.6 86.1 ± 0.1 30.4 ± 1.1 30.3 ± 0.5 30.6 ± 0.5 10−210^-2 75.3 ± 3.1 79.7 ± 0.9 81.9 ± 0.4 27.1 ± 0.7 27.6 ± 0.1 28.0 ± 0.1 10−110^-1 65.0 ± 3.3 72.6 ± 0.9 74.3 ± 0.5 26.3 ± 0.3 26.2 ± 0.2 26.2 ± 0.3 Comments. After four epochs, specialists average 8989-91%91\% accuracy on their own tasks but show little transfer to the other tasks. Task Arithmetic and TIES-Merging outperform weight averaging, with the best evaluated merges reaching 61.361.3, 71.871.8, and 72.1%72.1\% for the 130M, 370M, and 780M checkpoints, respectively, compared with 37.437.4, 50.450.4, and 51.7%51.7\% for averaging (Table 29). Larger checkpoints narrow the gap to specialists, although substantial deficits remain, especially on DBpedia and TREC (Table 30). The sweeps favor Task Arithmetic scales of λ=0.7λ=0.7-11, rather than the averaging scale λ=0.25λ=0.25, and sparse, rescaled TIES updates, with (λ,d)=(1.3,0.1)(λ,d)=(1.3,0.1) performing best among the four-epoch settings evaluated (Tables 31 and 32). Additionally, fine-tuning with the proximal penalty μ2‖θ−θ0‖F2 μ2\|θ-θ^0\|_F^2 reduces four-epoch averaging accuracy from 37.4%37.4\% at μ=0μ=0 to 26.2%26.2\% at μ=10−1μ=10^-1, near the pretrained baseline (Table 33). Relation to the analysis. These results are consistent with the trade-off in the local loss control (466) between the linear alignment term ⟨∇f(θ0),θmerge−θ0⟩F ∇ f(θ^0), _merge-θ^0 _F and the quadratic smoothness remainder. Weight averaging preserves neighborhood containment but rescales each task vector by 1/M1/M. For the 780M model parameters, ‖∑m‖F2/∑m‖m‖F2=1.04\| _mT_m\|_F^2/ _m\|T_m\|_F^2=1.04-1.071.07, indicating little net alignment among the task vectors. Together with the scale sweeps, this suggests that averaging suppresses useful task-specific adaptation: larger displacements can improve accuracy despite weaker containment guarantees, illustrating why containment alone does not ensure predictive performance (Appendix F.1). The proximal results suggest a similar trade-off. Limiting specialist drift can also suppress useful adaptation, and, unlike iterative federated learning, one-shot merging has no subsequent rounds to recover it. Finally, the uniform smoothness bound does not explain the larger accuracy gaps on DBpedia and TREC, motivating further investigation of the sensitivity-aware per-block scales discussed in Appendix F.2.

Instant research discovery