Distributed Learning with Selective State Space Models: Architecture-Aware Convergence Analysis
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.
Related
- Linear convergence rate analysis of Proximal-Tracking for nonsmooth distributed optimization
- Optimal Design of Distributed Inexact Gradient Tracking: Exact Worst-case Convergence Rate and Explicit Parameters
- Differentially Private Distributed Optimization via the Intrinsic Randomness of Zeroth-Order Estimators: Sampling Geometry, Convergence, and Privacy Guarantees
- Distributed Cross-Layer Optimization for Covert Multi-Hop, Multi-Modal Networks: Exponentially Fast Convergence and Robust Tracking
- Distributed Selective Inference for Quantile Regression
BibTeX
@misc{piaseczny2026distributed,
title = {{Distributed Learning with Selective State Space Models: Architecture-Aware Convergence Analysis}},
author = {Adam Piaseczny and Md Kamran Chowdhury Shisher and Shiqiang Wang and Christopher G. Brinton},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.02659},
eprint = {2610.02659},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.02659},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Distributed Learning with Selective State Space Models: Architecture-Aware Convergence Analysis
AU - Adam Piaseczny
AU - Md Kamran Chowdhury Shisher
AU - Shiqiang Wang
AU - Christopher G. Brinton
PY - 2026
DA - 2026/10/06
PB - arXiv
DO - 10.48550/arxiv.2610.02659
UR - https://arxiv.org/abs/2610.02659
AB - 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.
ER -