Title: Mitigating Classifier Dimensional Collapse via Random Masking in Federated Linear Probing
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: federated learning, data heterogeneity, dimensional collapse, pre-trained model

Abstract: Integrating a pre-trained model into federated learning (FL) is an emergent direction to facilitate the industrial deployment. Federated linear probing (FLP) is a practical paradigm that gains the communication efficiency from FL and generalization from pre-trained model, while still suffering from dimensional collapse that undermines its effectiveness. Dimensional collapse, originating from data heterogeneity, challenges the embedding space construction in FL, leading to suboptimal convergence and generalization. With a generalized frozen embedding extractor, FLP seems to be robust against dimensional collapse. However, in this paper, we emphasize that the dimensional collapse can also be represented in classifier construction, affecting the performance of the model. We propose FedRM to solve this problem, which randomly masks the dimension of the embedding and the classifier during the training to enforce the classifier to focus fairly on each dimension, guaranteeing diversity during decision generation. The simplicity of the method retains the communication efficiency of the FLP. We conduct empirical experiments to comprehensively evaluate the performance of FedRM. The results show FedRM achieves an overwhelming trade-off between efficiency and utility.