Research map / Federated learning and privacy / Personalized and heterogeneous FL
Improved Modelling of Federated Datasets using Mixtures-of-Dirichlet-Multinomials
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a Mixture-of-Dirichlet-Multinomials model, fit with a fully federated generalized-EM algorithm, to learn the distribution of true federated clients' label histograms so that centralized proxy data can be partitioned into realistic simulated clients for server-side FL simulation.
Problem
Server-side proxy data is usually a single centralized dataset without client partitions, and naive IID partitioning fails to reflect the statistical heterogeneity of true federated clients, so hyperparameter tuning simulations are unrepresentative.
Approach
Model each client's categorical feature histogram (and sample count) as a draw from a mixture of Dirichlet-multinomials, infer its parameters by MLE using only aggregated client statistics via generalized EM with closed-form and fixed-point updates, then sample histograms to subsample centralized data into simulated clients.
Key results
- On FEMNIST, mean absolute accuracy difference between true clients and MDM simulated clients was 1.6% versus 3.3% for fully IID simulation.
- On CIFAR10, mean absolute accuracy difference was 0.8% for MDM simulation versus 6.6% for fully IID simulation.
- Inferred MDM parameters converge over rounds in synthetic experiments, with normalized MSE decreasing across low, medium and high heterogeneity settings.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use learned MDM simulated clients instead of fully IID simulated clients | FedAvg training simulations · FEMNIST | mean absolute accuracy difference from true-client training: 3.3% -> 1.6% | Helped |
| Use learned MDM simulated clients instead of fully IID simulated clients | FedAvg training simulations · CIFAR10 | mean absolute accuracy difference from true-client training: 6.6% -> 0.8% | Helped |
| Use learned MDM simulated clients instead of fully IID simulated clients | Simulating client histograms · FEMNIST, Folktables race, and Folktables income | client-distribution similarity | Helped |
| Use 3 mixture components instead of 2 (ablation) | MDM client simulation · FEMNIST | coverage of client-histogram clusters | Helped |
| Use K≥3 instead of K=1 (ablation) | MDM client simulation · Folktables race | client-distribution similarity | Helped |
| Use K>1 instead of K=1 (ablation) | MDM client simulation · Folktables binned income | client-distribution similarity | Helped |
| Use a larger number of mixture components (ablation) | MDM client simulation · Folktables binned income | client-distribution similarity | Helped |
Surprising findings
- Fully IID simulated clients give the highest accuracies across nearly all hyperparameter settings, i.e. an overly optimistic prediction of true-client performance.
- On FEMNIST, MDM simulations tend to overestimate performance on true clients, leaving a non-negligible gap, which the authors attribute to heterogeneity in features beyond label distribution.
Limitations
- Non-IID effects beyond the modelled categorical feature (e.g. feature distribution heterogeneity) are not captured, leaving residual gaps on FEMNIST.
- Only one categorical feature is modelled at a time; the number of mixture components K must be chosen by validation likelihood.
- MDM simulations overestimate performance on true clients on FEMNIST.
Experimental setup
Datasets: CIFAR10 FEMNIST Folktables synthetic MDM data
Benchmarks: CIFAR10 FEMNIST Folktables
Metrics: normalized MSE of inferred parameters test accuracy gap between simulated and true clients mean absolute accuracy difference log likelihood for choosing K
Baselines: fully IID simulation conditionally IID simulation
Builds on: Dirichlet-multinomial distribution FedAvg pfl-research
Closest papers by meaning
- Convergence Analysis of Split Federated Learning on Heterogeneous Data (NeurIPS 2024 · similarity 0.88)
- Federated Learning with Sample-level Client Drift Mitigation (AAAI 2025 · similarity 0.87)
- Heterogeneity-Guided Client Sampling: Towards Fast and Efficient Non-IID Federated Learning (NeurIPS 2024 · similarity 0.87)
- Adaptive Federated Learning with Auto-Tuned Clients (ICLR 2024 · similarity 0.87)
- FedImpro: Measuring and Improving Client Update in Federated Learning (ICLR 2024 · similarity 0.87)
- Streaming Federated Learning with Markovian Data (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.