Mitigating Catastrophic Forgetting in Personalized Federated Learning for Edge Devices Using State-Space Models

Edge computing involves distributed devices operating in dynamic environments with diverse resource constraints and highly heterogeneous data distributions. In these settings, personalized federated learning (PFL) provides a collaborative learning framework that preserves the unique data characteristics of each device. However, PFL models often encounter <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">bidirectional catastrophic forgetting</i>. During consecutive training rounds, the personalized characteristics learned by local models are readily overwritten by updates from the global model, while the global model's shared representations are degraded by distribution shifts arising from heterogeneous local data. This challenge is further amplified in edge computing settings, where devices must adapt to fast-changing heterogeneous data. To address this issue, we propose FedSSM, a novel framework that leverages <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">state-space models</i> to mitigate forgetting in PFL. By capturing the temporal evolution of local model parameters through hidden states, the framework enhances the retention of critical knowledge throughout training rounds. Extensive experiments on multiple benchmark datasets demonstrate that FedSSM outperforms various state-of-the-art PFL algorithms, particularly with high data heterogeneity. The code can be found at: <uri xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">https://github.com/wei-d-zhang/FedSSM</uri>.