Research map / Federated learning and privacy / Personalized and heterogeneous FL
Uncertainty-Based Extensible Codebook for Discrete Federated Learning in Heterogeneous Data Silos
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 UEFL, an iterative federated learning framework that uses an uncertainty-triggered, K-means-initialized extensible VQ-style codebook (MC Dropout entropy) to handle feature-heterogeneous data silos.
Problem
Federated learning across heterogeneous data silos with divergent distributions yields a global model with lower accuracy and higher uncertainty than local models, and existing discrete-representation FL methods generalize poorly to large distribution shifts.
Approach
Use MC Dropout predictive entropy to identify silos with high uncertainty, then add new codewords (K-means initialized from the trained encoder's features) visible only to those silos, iteratively extending a shared discrete codebook.
Key results
- On rotated multi-domain data, UEFL reaches 0.920 mean accuracy on MNIST vs 0.815 for DisTrans and 0.780 for FedAvg, with entropy reduced from 0.273 to 0.149.
- On PACS domain generalization, UEFL improves average accuracy from 0.803 to 0.8453, above FedSR (0.834) and FedIIR (0.837).
- Computation overhead is small: CIFAR10 parameters rise from 14.991M to 15.491M (3.34%) and GPU runtime from 16.154 ms to 16.733 ms.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use UEFL instead of DisTrans | Federated learning with heterogeneous silos · GTSRB | mA: 0.898 -> 0.942 | Helped |
| Use UEFL instead of FedAvg | Federated learning with heterogeneous silos · MNIST, FMNIST, GTSRB, CIFAR10, CIFAR100 | mE: 0.273 -> 0.149; 0.273 -> 0.167; 0.636 -> 0.0239; 0.177 -> 0.0222; 1.91 -> 0.655 | Helped |
| Add feature discretization (ablation) | VQ-FedAvg · MNIST | mean accuracy: 0.834 -> 0.907 | Helped |
| Gradually extend a small codebook instead of starting with a static large codebook (ablation) | UEFL · CIFAR100 | accuracy: 0.13 to 0.34 | Helped |
| Use K-means to initialize new codewords (ablation) | UEFL · MNIST | mA: 0.888 -> 0.945; 0.554 -> 0.906; 0.622 -> 0.929 | Helped |
| Increase initial codewords from 32 to 64 (ablation) | UEFL · GTSRB | mA: 0.956 -> 0.965; 0.966 -> 0.974; 0.911 -> 0.926 | Mixed |
| Use 4 codeword segments instead of fewer segments (ablation) | UEFL · CIFAR100 | performance | Helped |
| Lower the uncertainty threshold (ablation) | UEFL · heterogeneous federated learning | performance | Helped |
| Replace Monte Carlo Dropout with Deep Ensembles (ablation) | UEFL uncertainty evaluation · MNIST, FMNIST, GTSRB, CIFAR10, CIFAR100 | mA: 0.920 -> 0.926; 0.850 -> 0.853; 0.942 -> 0.940; 0.720 -> 0.717; 0.326 -> 0.331 | Mixed |
| Replace Monte Carlo Dropout with Deep Ensembles (ablation) | UEFL uncertainty evaluation · MNIST, FMNIST, GTSRB, CIFAR10, CIFAR100 | mE: 0.149 -> 0.211; 0.167 -> 0.289; 0.0239 -> 0.041; 0.0222 -> 0.043; 0.655 -> 0.873 | Hurt |
| Use Deep Ensembles instead of Monte Carlo Dropout (ablation) | UEFL uncertainty evaluation · uncertainty evaluation | computational resources | Hurt |
| Add codewords after round 60 (ablation) | UEFL · CIFAR10 | accuracy and uncertainty | No clear effect |
Surprising findings
- Deep Ensembles give accuracy similar to MC Dropout but substantially higher uncertainty, and cost far more compute.
- Neural collapse training, which removes dropout, still underperforms UEFL on CIFAR-100 (0.3482 vs 0.5074 reported in appendix table).
Limitations
- Uncertainty threshold gamma and codebook extension count are hyperparameters tuned per dataset, and the paper notes a runtime/performance trade-off at low thresholds.
- Evaluated only on image classification with a small number of silos (9 main experiments; 50 and 100 clients in additional settings), and mostly with rotation-based feature shifts.
- Theoretical analysis (discretization generalization bounds) is loosely connected to the empirical method.
Experimental setup
Datasets: MNIST FMNIST GTSRB CIFAR10 CIFAR100 Rotated MNIST PACS EMNIST-L
Benchmarks: MNIST FMNIST GTSRB CIFAR10 CIFAR100 Rotated MNIST PACS EMNIST-L
Metrics: mean accuracy mean entropy (uncertainty) mean perplexity
Baselines: FedAvg DisTrans FedSR FedIIR FedCR VHL FedBR Deep Ensembles
Builds on: FedAvg VQ-VAE DisTrans Monte Carlo Dropout
Closest papers by meaning
- FedRSPO+: A Heterogeneity-aware Algorithm for Decision-focused Federated Learning (NeurIPS 2026 · similarity 0.89)
- A Bayesian Approach for Personalized Federated Learning in Heterogeneous Settings (NeurIPS 2024 · similarity 0.88)
- Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework (ICML 2025 · similarity 0.88)
- FedGPS: Statistical Rectification Against Data Heterogeneity in Federated Learning (NeurIPS 2025 · similarity 0.88)
- Bridging Model Heterogeneity in Federated Learning via Uncertainty-based Asymmetrical Reciprocity Learning (ICML 2024 · similarity 0.87)
- FedUV: Uniformity and Variance for Heterogeneous Federated Learning (CVPR 2024 · 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.