FedRIR: Rethinking Information Representation in Federated Learning

TL;DR

FedRIR tackles the dual challenge of achieving global generalization and local personalization in non-IID federated learning. It introduces Masked Client-Specific Learning ($\mathrm{MCSL}$) to isolate client-specific features and an Information Distillation Module ($\mathrm{IDM}$) to distill pure global representations using the variational CLUB bound, yielding $f_p^i=[f_{g}, f_{cs}^{i}]$ for customization. Across pathological, practical, and real-world heterogeneity, FedRIR consistently outperforms state-of-the-art pFL baselines in accuracy while maintaining comparable communication overhead, and ablations confirm the necessity of both MCSL and IDM. The approach scales to large numbers of clients and remains robust under varying participation, making it a practical framework for edge-enabled FL where data heterogeneity and privacy are pivotal concerns.

Abstract

Mobile and Web-of-Things (WoT) devices at the network edge generate vast amounts of data for machine learning applications, yet privacy concerns hinder centralized model training. Federated Learning (FL) allows clients (devices) to collaboratively train a shared model coordinated by a central server without transfer private data, but inherent statistical heterogeneity among clients presents challenges, often leading to a dilemma between clients' needs for personalized local models and the server's goal of building a generalized global model. Existing FL methods typically prioritize either global generalization or local personalization, resulting in a trade-off between these two objectives and limiting the full potential of diverse client data. To address this challenge, we propose a novel framework that simultaneously enhances global generalization and local personalization by Rethinking Information Representation in the Federated learning process (FedRIR). Specifically, we introduce Masked Client-Specific Learning (MCSL), which isolates and extracts fine-grained client-specific features tailored to each client's unique data characteristics, thereby enhancing personalization. Concurrently, the Information Distillation Module (IDM) refines the global shared features by filtering out redundant client-specific information, resulting in a purer and more robust global representation that enhances generalization. By integrating the refined global features with the isolated client-specific features, we construct enriched representations that effectively capture both global patterns and local nuances, thereby improving the performance of downstream tasks on the client. The code is available at https://github.com/Deep-Imaging-Group/FedRIR.

Figures (5)

• Figure 1: The workflow of the proposed FedRIR framework.
• Figure 2: The data distribution on FMNIST and Cifar10 datasets in pathological and practical settings with 20 clients. The size of a circle indicates the number of local data.
• Figure 3: t-SNE visualization of global features extracted by FedRoD, FedProto, FedCP and FedRIR on Cifar10 dataset, with color represents a different class.
• Figure 4: The impact of different mask ratios on the performance of the FedRIR method.
• Figure 5: t-SNE visualization of client-specific features extracted by FedRIR w/o MCSL and FedRIR on OfficeCaltech10 dataset, with color represents a different client.