# Uncertainty-Based Extensible Codebook for Discrete Federated Learning in Heterogeneous Data Silos

arXiv: https://arxiv.org/abs/2402.18888 · HTML: https://mlsanity.com/paper/2402.18888/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 43 references that are in the map.
- First posted on arXiv 29 Feb 2024.

## Builds on (most cited first, 30 of 31)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Differentially Private Federated Learning: A Client Level Perspective](https://mlsanity.com/paper/1712.07557/) (2017)
- [Domain Generalization for Object Recognition with Multi-task Autoencoders](https://mlsanity.com/paper/1508.07680/) (2015)
- [Multiplicative Normalizing Flows for Variational Bayesian Neural Networks](https://mlsanity.com/paper/1703.01961/) (2017)
- [Federated Meta-Learning with Fast Convergence and Efficient Communication](https://mlsanity.com/paper/1802.07876/) (2018)
- [RSA: Byzantine-Robust Stochastic Aggregation Methods for Distributed Learning from Heterogeneous Datasets](https://mlsanity.com/paper/1811.03761/) (2018)
- [Uncertainty-based Continual Learning with Adaptive Regularization](https://mlsanity.com/paper/1905.11614/) (2019)
- [The Distributed Discrete Gaussian Mechanism for Federated Learning with Secure Aggregation](https://mlsanity.com/paper/2102.06387/) (2021)
- [DEUP: Direct Epistemic Uncertainty Prediction](https://mlsanity.com/paper/2102.08501/) (2021)
- [Uncertainty Baselines: Benchmarks for Uncertainty & Robustness in Deep Learning](https://mlsanity.com/paper/2106.04015/) (2021)
- [The Skellam Mechanism for Differentially Private Federated Learning](https://mlsanity.com/paper/2110.04995/) (2021)
- [Decentralized Federated Learning through Proxy Model Sharing](https://mlsanity.com/paper/2111.11343/) (2021)
- [Virtual Homogeneity Learning: Defending against Data Heterogeneity in Federated Learning](https://mlsanity.com/paper/2206.02465/) (2022)
- [FedFA: Federated Learning with Feature Anchors to Align Features and Classifiers for Heterogeneous Data](https://mlsanity.com/paper/2211.09299/) (2022)
- [Analyzing the Role of Model Uncertainty for Electronic Health Records](https://mlsanity.com/paper/1906.03842/) (2019)
- [FedBR: Improving Federated Learning on Heterogeneous Data via Local Learning Bias Reduction](https://mlsanity.com/paper/2205.13462/) (2022)
- [Addressing Heterogeneity in Federated Learning via Distributional Transformation](https://mlsanity.com/paper/2210.15025/) (2022)
- [PMFL: Partial Meta-Federated Learning for heterogeneous tasks and its applications on real-world medical records](https://mlsanity.com/paper/2112.05321/) (2021)

## Built on by (most cited first, 3 of 3)

- [LoRA-LiteE: A Computationally Efficient Framework for Chatbot Preference-Tuning](https://mlsanity.com/paper/2411.09947/) (2024)
- [Federated Learning-Based Data Collaboration Method for Enhancing Edge Cloud AI System Security Using Large Language Models](https://mlsanity.com/paper/2506.18087/) (2025)
- [FERA: Uncertainty-Aware Federated Reasoning for Large Language Models](https://mlsanity.com/paper/2605.10082/) (2026)

## Abstract (arXiv)

Federated learning (FL), aimed at leveraging vast distributed datasets, confronts a crucial challenge: the heterogeneity of data across different silos. While previous studies have explored discrete representations to enhance model generalization across minor distributional shifts, these approaches often struggle to adapt to new data silos with significantly divergent distributions. In response, we have identified that models derived from FL exhibit markedly increased uncertainty when applied to data silos with unfamiliar distributions. Consequently, we propose an innovative yet straightforward iterative framework, termed Uncertainty-Based Extensible-Codebook Federated Learning (UEFL). This framework dynamically maps latent features to trainable discrete vectors, assesses the uncertainty, and specifically extends the discretization dictionary or codebook for silos exhibiting high uncertainty. Our approach aims to simultaneously enhance accuracy and reduce uncertainty by explicitly addressing the diversity of data distributions, all while maintaining minimal computational overhead in environments characterized by heterogeneous data silos. Extensive experiments across multiple datasets demonstrate that UEFL outperforms state-of-the-art methods, achieving significant improvements in accuracy (by 3%--22.1%) and uncertainty reduction (by 38.83%--96.24%). The source code is available at https://github.com/destiny301/uefl.