Distribution-Aware Personalised Aggregation for Federated Deep Learning Under Non-IID Data

Abstract

Federated learning enables collaborative deep neural network training across distributed clients without sharing raw data, yet its performance degrades substantially when local data distributions are non-identically distributed (non-IID). Existing aggregation strategies either treat all clients uniformly or require expensive bi-level optimisation, failing to explicitly leverage the structural similarity among client data distributions. We propose DAPA (Distribution-Aware Personalised Aggregation), a personalised federated deep learning framework that adaptively tailors the global aggregation to each client based on inter-client distribution similarity. DAPA operates in three stages. In the sketch stage, each client computes a lightweight distribution sketch that summarises its local label and feature statistics through class-conditional moment vectors extracted from the deep network’s penultimate layer. In the affinity stage, the server constructs a pairwise client affinity matrix from these sketches using an approximate Wasserstein distance and derives personalised aggregation weight vectors via a softmax-temperature mechanism. In the clustering stage, a hierarchical two-level aggregation combines models within automatically discovered client clusters and then blends across clusters with adaptive mixing coefficients. To safeguard privacy, the distribution sketches are protected with a calibrated Gaussian mechanism that satisfies Rényi differential privacy. Theoretical analysis establishes a convergence bound showing that DAPA achieves a tighter error floor than uniform averaging under distribution heterogeneity. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny-ImageNet under Dirichlet-controlled non-IID partitions demonstrate that DAPA outperforms nine state-of-the-art baselines, improving average test accuracy by 2.4 to 6.8 percentage points while maintaining competitive communication efficiency. Ablation studies confirm the contribution of each component, and privacy analysis verifies that the accuracy gain persists under strict differential privacy budgets. These findings advance the application of deep learning in privacy-sensitive distributed environments and offer a principled data-mining-based approach to characterising client heterogeneity in federated systems.