Exploring personalized federated learning from a distribution-based perspective

Personalized federated learning (PFL) is a promising technique for tackling data heterogeneity in federated learning systems. Recently, Bayesian neural networks (BNNs) have been introduced into the PFL framework to enable uncertainty quantification and improve performance in data-scarce settings. Despite these advantages, existing BNN-based PFL methods face two key challenges in practical applications. First, in real-world scenarios, client heterogeneity often arises in the form of group-wise variation, which cannot be adequately captured by a single shared distribution as assumed in prior work. Second, existing methods rely on deterministic or stochastic approximation techniques for posterior inference, which lead to substantial computational and memory overhead, hindering their scalability and deployment. To address these limitations, we propose DBFed, a novel BNN-based PFL framework from a distribution-based perspective. DBFed introduces group-specific distributions to better model the structural heterogeneity commonly observed in federated settings. Moreover, DBFed employs a rank-1 parameterization technique to map uncertainty from the weight space to a low-dimensional subspace, significantly reducing the computational and memory overhead. Theoretically, we establish the effectiveness of the rank-1 parameterization approach. Empirically, extensive experiments on diverse datasets demonstrate that DBFed consistently outperforms alternative PFL baselines in a heterogeneous setting.