选择性SSM分布式训练的架构相关收敛分析

Mamba2等现代状态空间模型（SSM）将线性时间序列建模与循环状态空间动力学结合，为Transformer提供有吸引力的替代方案。但SSM在分布式学习环境中的行为仍不清楚。现有标准联邦学习方法大多不区分架构，没有考虑现代选择性SSM的稳定性、选择性和状态空间参数化。为此，我们推导单层和多层选择性SSM的架构相关梯度与平滑性界，以及FedAvg和FedProx的收敛界，刻画循环稳定性、输入相关离散化和状态投影范数如何影响联邦优化。随后，我们在教师SSM生成的序列上，以遵循所分析递推的学习器数值验证单层界。利用分析，我们形成关于本地训练和客户端异质性影响的预期，并通过比较Mamba2在六个文本域语言建模上的九种联邦学习算法检验这些预期。实验说明，SSM专属的界可为解释实际联邦算法行为提供基础。