Personalized PCA for Federated Heterogeneous Data

As the high dimensional data generation/storage shifts from data centers to millions of edge devices, PCA algorithms also need to adapt to federated systems to reveal insights about the distributed data. One of the prominent challenges in Federated Learning (FL) is that each edge device has a limited number of samples, and therefore collaboration among clients is necessary for learning tasks. Another challenge is heterogeneous distribution of data across devices, which necessitates careful design of algorithms that enable collaboration of devices with different data distributions. While many such federated supervised learning algorithms were proposed in recent years, heterogeneity for unsupervised FL algorithms (such as PCA) has received less attention. In this work, our goal is to enable collaborations of heterogeneous clients in learning personalized Principal Components (PCs). To this end, we develop a hierarchical Bayesian framework for discovering individual PCs; and inspired by this, we formulate an optimization problem related to maximum likelihood estimation of the PCs. To solve the optimization problem, we propose an alternating Stiefel gradient descent algorithm. Analytically, we prove the convergence result for our proposed algorithm; and empirically, we show that our method outperforms local and global estimation of PCs in various heterogeneous settings in terms of the reconstruction error.