Consensus Memory Alignment for Effective Knowledge Transfer in Personalized Federated Learning

Abstract

Personalized Federated Learning (PFL) addresses performance degradation in non-IID settings by preserving personalized local models. However, the inherent reinitialization of local models in each communication round weakens cross-client knowledge transfer, which disrupts the balance between global and local knowledge. This imbalance restricts the utilization of shared representations, resulting in biased feature extraction that ultimately limits the model’s personalization performance. To solve this problem, we propose Feature Distribution Alignment (FDA), a general FL framework that enhances collaborative knowledge transfer. It consists of two components, namely Consensus Knowledge Guidance (CKG) and Local Knowledge Memory (LKM). CKG employs a domain adversarial approach adapted to FL settings, guiding local feature extractors toward globally consistent representations. LKM retains historical local features via attention-based fusion and mutual information maximization, enhancing feature alignment while reducing redundancy. This combined mechanism ensures the persistent retention of prior knowledge in local models despite continuous global model updates. By integrating these components, the FDA enables local feature extractors to better align with global knowledge while maintaining personalization. This design effectively balances global and local knowledge while improving local model performance. Theoretical analysis provides generalization guarantees supporting the effectiveness of FDA in facilitating server-to-client knowledge transfer. Extensive experiments on four datasets under three statistically heterogeneous settings show that FedAvg+FDA outperforms leading PFL methods in terms of effectiveness, stability, and fairness. Our code is available at https://github.com/NiantaoZ/FedFDA.