Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning (2210.00226v5)

Abstract: Federated learning aims to train models collaboratively across different clients without the sharing of data for privacy considerations. However, one major challenge for this learning paradigm is the {\em data heterogeneity} problem, which refers to the discrepancies between the local data distributions among various clients. To tackle this problem, we first study how data heterogeneity affects the representations of the globally aggregated models. Interestingly, we find that heterogeneous data results in the global model suffering from severe {\em dimensional collapse}, in which representations tend to reside in a lower-dimensional space instead of the ambient space. Moreover, we observe a similar phenomenon on models locally trained on each client and deduce that the dimensional collapse on the global model is inherited from local models. In addition, we theoretically analyze the gradient flow dynamics to shed light on how data heterogeneity result in dimensional collapse for local models. To remedy this problem caused by the data heterogeneity, we propose {\sc FedDecorr}, a novel method that can effectively mitigate dimensional collapse in federated learning. Specifically, {\sc FedDecorr} applies a regularization term during local training that encourages different dimensions of representations to be uncorrelated. {\sc FedDecorr}, which is implementation-friendly and computationally-efficient, yields consistent improvements over baselines on standard benchmark datasets. Code: https://github.com/bytedance/FedDecorr.

• The paper identifies that data heterogeneity in federated learning exacerbates dimensional collapse of model representations by inducing low-rank bias during local training.
• It introduces FedDecorr, a novel method using a computationally efficient regularization term during local training to decorrelate representation dimensions and mitigate collapse.
• Experiments show FedDecorr consistently improves performance over baseline methods, especially under high data heterogeneity, enhancing federated learning applicability in diverse real-world settings.

Dimensional Collapse in Federated Learning: Mitigation via FedDecorr

This paper addresses a critical issue in Federated Learning (FL), particularly under conditions of data heterogeneity: dimensional collapse of representations. FL is a decentralized model training paradigm designed to maintain data privacy by training models across multiple clients, each holding local data, without centralizing the data. However, models trained through FL often face the challenge of data heterogeneity, where local data varies significantly across clients. This discrepancy in data distribution can lead to the model's performance degradation, a phenomenon explored in this work through the concept of dimensional collapse.

Dimensional collapse refers to the situation where representations produced by the model occupy a lower-dimensional space than intended. The paper empirically demonstrates that heterogeneity across client data exacerbates this collapse, impacting both local and global models aggregated in FL schemes. By analyzing gradient dynamics, the authors establish that data heterogeneity fosters low-rank bias in the weight matrices during local training, further leading to the dimensional collapse observed in model representations.

The central contribution of the paper is the introduction of FedDecorr, a novel technique designed to counteract dimensional collapse. This method deploys a regularization term during local training to ensure different dimensions of representations remain uncorrelated. The regularization is computationally efficient as it works by minimizing the Frobenius norm of the representation’s correlation matrix.

Through comprehensive experiments using standard benchmarks such as CIFAR10, CIFAR100, and TinyImageNet, FedDecorr consistently outperformed baseline FL methods across varying degrees of data heterogeneity. Notable improvements were seen especially under settings with high heterogeneity, suggesting that FedDecorr is effective where traditional FL methods falter. The paper also performs ablation studies demonstrating the robustness of FedDecorr across different numbers of clients and attention to the local training epochs.

The implications of this research are significant: FedDecorr not only enhances the performance of existing FL methods but also points toward optimization approaches that integrate decorrelation to mitigate the effects of local data heterogeneity. The findings offer potential applications for deploying federated learning in real-world scenarios involving diverse data distributions, such as in personal healthcare devices, mobile applications, and more, where client data is inherently heterogeneous.

Future research might explore extending FedDecorr principles to other machine learning paradigms impacted by data heterogeneity and dimensional collapse. The technique's adaptability to non-linear neural network architectures or its integration with personalized federated learning frameworks could further broaden its application landscape. FedDecorr establishes a path for efficiently handling intrinsic data challenges in decentralized learning systems, providing a crucial tool for privacy-preserving AI applications.