Uncertainty-Based Extensible Codebook for Discrete Federated Learning in Heterogeneous Data Silos
Authors: Tianyi Zhang
Yu Cao
Dianbo Liu
Venue: ICML 2025 poster (2025)
Keywords: federated learning
heterogeneous data
discretization
Primary area: social_aspects->fairness

Abstract: 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 \emph{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.

Uncertainty-Bast -Based Extensible Codebook for Discrete Federated Learning in Heterogencous Data Silos
Tianyi Zhang'
Yu Cao
Dianbo Liu
Abstract
FL
UEFL (ours)
Federated learning (FL), aimed at leveraging vast distributed datasets, confronts crucial chal- lenge: the heterogeneity of data across different silos. While previous studies have explored dis- crete representations to enhance model general- ization across minor distributional shifts, these approaches often struggle 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 distri- butions. Consequently, we propose an innovative yet straightforward iterative framework, termed Uncertainty- Based Extensible-Codebook Feder- ated Learning (UEFL.). This framework dynami- cally maps latent features to trainable discrete vec- tors, assesses the uncertainty, and specifically ex- tends the discretization dictionary or codebook for silos exhibiting high uncertainty. Our approach aims to simultaneously enhance accuracy and re- duce uncertainty by explicitly addressing the di- versity of data distributions, all while maintain- ing minimal computational overhead in environ- ments 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 re- duction (by 98.83%-96.244% The source code available at https://gituub..om/dssiii
Local Acturacy/ Acturacy/ Uncertalntytt.a.o.ooom 0.924/0.102
Encodling Discretization
Encoding
Encoding
Global Accuracy/Uncartainty 0.803/0.21
Global Accuracy/Unceertaiity 0.845/0.149
In the case of heterogeneous data silos, the global model of regular FL performs poorly compared to local mod- els. By discretizing different domains into distinct latent spaces, our UEFL improves both accuracy and uncertainty. The reported values in the figure represent the average accu- racy and uncertainty across the various data silos.
Figure 1:
milos-while ensuring privacy, has become increasingly cru- cial in the digital era, particularly with the explosion of data from mobile sources. Despite its pivotal role in distributed computing, FL confronts formidable challenge: the hetero- geneity of data across different silos. Such diversity often results in a significant performance gap when integrating up- dates from local models into the global model. In Figure 1, compare the mean accuracy of local FL models with that of the global model after integration when addressing data silos with different distributions. While local models may perform impressively within their own data domains, the aggregated global model often struggles to achieve similar performance levels after synthesizing updates from these varied data sources. This issue is especially pronounced in FL due to its reliance varied data sources.
Recent studies (Ghosh et al., 2020; Agarwal et al., 2021; Liu et al., 2021; Kairouz et al., 2021a; Zhang et al., 2022; Yuan et al., 2022) have made significant advancements address- ing data heterogeneity within FL, with one notable approach being the use of discrete representations to enhance model robustness against minor data shifts. Nonetheless, this strat- egy struggles to generalize models to data silos exhibiting significant distributional differences. Furthermore, these methods face difficulties in adapting to unseen data distri- butions, as they typically require the entire model to be re-traine re-trained. Such constraints limit their flexibility in adapt- ing to the dynamically changing data landscapes, posing
1. Introduction
Federated Learning (FI.), well known for its capacity to har- ness data from diverse devices and locations termmd data
'University of Minnesota, MN, USA NNational University of Singapore, Singapore. Correspondence to: Tianyi Zhang <zhan9167 @umn.emumu.... Dianbo Liu <dianbo@ nus eds.edu.sg>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Envertainty-Basell Extensible-Codcbook Federated Learning (UEFL)
challenges for their applicability in real-world scenarios.
enabling rapid convergence during codebook training.
Moreover, we identify another critical issue impacting the model's performance across diverse data silos: increased uncertainty. as shown in Figure 1. The global model's 's ac- curacy not only deteriorates, but its uncertainty also trends upwards, signaling increased prediction instability. To ad- dress these challenges, we introduce Uncertainty-Basee Extensible-codebook Federated Learning (UEFL), novel methodology that explicitly distinguishes between data dis- tributions to improve both accuracy and uncertainty.
We propose a novel data-driven FL framework, named Uncertainty-Based Extensible-codehook Federated Learning (UEFL), which merges the extensible code- book with an uncertainty evaluator. This framework iteratively identifies data from diverse distributions by assessing uncertainty without requiring direct data ac- cess. It then processes this data by initializing new codewords to complement the existing codebook, en- suring that cach iteration focuses on training the ex- pandable codebook, thus allowing UEFL to adapt seam- lessly to new data distributions.
Specifically, our design features an advanced codebook com- prising a predetermined number of latent vectors (i.e. code- words), and employs a discretizer to assign encoded image features their closest codewords. These codewords, acting as latent representations, are passed subsequent layers for processing. The codewords are dynamically trained to align with the latent features generated by the image encoder. To mitigate performance degradation when integrating local models from data silos with varying distributions, we initial- ize a small, shared codebook for all clients. Additional spe- cific codewords are then introduced for individual client use, ensuring explicit differentiation between them. Since the initial codebook is small and requires only few extensions, the final size remains compact, minimizing the associated computational overhead. Given the privacy constraints in FL., which restrict direct data access, we incorporate an un- certainty evaluator using Monte Carlo Dropout. This evalua- tor identifies data from diverse distributions, marked by high uncertainty. During training, UEFL systematically distin- guishes between these varied distributions and dynamically adds new codewords to the codebook until all distributions are sufficiently represented. In the initial training cycle, shared codewords are randomly initialized. However, in subsequent cycles, the fully trained image encoder is lever- aged to initialize new codewords using K-means, aligning them more closely with the data distribution and facilitating faster adaptation various distributions. As a result, our UEFL model can accommodate data from previously unseen distributions with fewer communication rounds, making it applicable for enhancing other FL algorithms.
Our empirical evaluation across various datasets demonstrates that our approach significantly reduces uncertainty by 38.83%%-...22 96.2%% and enhances model accuracy by 3%-22. %, evidencing the effectiveness of UEFL in managing data heterogeneity in FL.
2. Related Work
2.1. Federated Learning
Federated learning (Konečný et al., 2016; Geyer et al., 2017: Chen et al., 2018; Hard et al., 2018; Yang et al., 2019; Ghosh et al., 2020) represents a cutting-edge distributed learning paradigm, specifically designed to exploit data and computational resources across edge devices. The Feder- ated Averaging (FedAvg) algorithm (McMahan et al., 2017), introduced to address the challenges of unbalanced and non- IID data, optimizes the trade-off between computation and communication costs by reducing the necessary communica- tion rounds for training deep networks. FL faces numerous statistical challenges, with data heterogeneity being one of the most critical. In real- woold applications, data collected across different clients often varies significantly in terms of distribution, feature space, and sample sizes.
Several methodologies (Zhao et al., 2018; Li et al., 2018; 2019; Kalra et al., 2023) have been developed to address this pivotal issue. PMFL (Zhang et al., 2022) approaches the heterogeneity challenge by drawing inspiration from meta-learning and continual learning, opting to integrate losses from local models over the aggregation of gradi- ents or parameters. DisTrams ( (Yuan et al., 2022) enhances FL performance through train and test-time distributional transformations, coupled with a novel double-input-channel model architecture. Mcanwhile, FCCL (Huang et al., 2022) employs knowledge distillation during local updates facil- itate the sharing of inter and intra domain insights without compromising privacy, and utilizes unlabeled public data to foster a generalizable representation amidst domain shifts. Additionally, the discrete approach to addressing hetero- geneity by (Liu et al., 2021). provide further inspiration and valuable perspectives for our research endeavors.
To summarize, our contributions are as follows:
We identify a significant increase in model uncertainty across silos with diverse data distributions in FL, high- lighting the challenge of data heterogeneity.
To address this heterogeneity, we introduce an exten- sible codebook approach that distinguishes between data distributions by stepwise mapping them to dis- tinct, trainable latent vectors (i.e. codewords). This methodology allows for efficient initialization of newly added codewords using a K m-mans algorithm, closely aligning with the training data feature distributions and
2

---

Uncertainty-Bane Extensible-Codcbook Federated Learning (UEFL)
2.2. Uncertainty
where 0 denots the global model parameters, Og is the kth local model parameters, Ok is the kth model gradients, ng is the number of samples for data silo k, and n is the total number of samples for all K silos.
Recently, the study of uncertainty modeling has gained significant prominence across various research fields, no- tably within the machine learning community (Chen al., 2014; Blundell et al., 2015; Kendall & Gal, 2017: Louizos & Welling, 2017: Lahlou et al., 2021; Nado et al., 2021; Gawlikowski et al., 2021). This surge in interest is driven by the critical need to understand and quantify the inherent ambiguity in complex datasets. Techniques such as Monte Carlo Dropout (Gal& & Ghahramani, 201 6). which introduces variability in model outputs through the use of dropout lay- ers, and Deep Ensembles (Lakshminarayanan et al., 2017), which leverages multiple models with randomly initialized weights trained on identical datasets to evaluate uncertainty, exemplify the advancements this area. Furthermore, the application of uncertainty modeling has extended beyond traditional domains, impacting fields such as healthcare (Dusenberry et al., 2020) and continual learning (Ahn et al., 2019).
At the end of each iteration, assessing uncertainty through Monte Carlo Dropout is essential, given the privacy con- straints of FL, which limit direct access to client data. By evaluating uncertainty against pre-established threshold, we identify data from heterogeneous distributions. When such data are detected, we augment the codebook with new codewords initialized by K-means. These newly gener- ated codewords are then exclusively accessible to the cor- responding heterogeneous clients, updating the codebook size for the kth client from eg to 1% + v, as described in Algorithm 1. This process leverages the fully adapted en- coder from previous iterations, utilizing K-means to ensure the new codewords are closely aligned with the actual data distribution, thereby facilitating faster convergence during training. Additionally, since the extended codewords are specific to individual client data and are not included in the integration with other local models, our method ensures that latent features from different distributions remain explicitly differentiated. Consequently, the global model performs bet- ter after integration, effectively handling data heterogencity.
3. Methodology
3.1. Overall Architecture
Figure 2 illustrates the workflow of our UEFL. Consider multiple data distributions D1, D2.... DM. with data sam- ples E # where H, W, and D denote the in- put image's 's height, width, and channel count, respectively, drawn from these M distributions. Upon distributing the global model to local clients, data samples undergo local encoding via a shared encoder 0g into feature represen- tations: = E Rgxxuxd, with h, w, and d representing the features' shape. Subsequently, these features are reshaped into vectors = gdxdd where / is the number of tokens, and divided into s segments 2, RRx e Vi, with indicating the segment count. Each segment is mapped to the closest code- word in the codebook via discretizer 0D. then reassembled into complete vectors for classification. The classifier Oc then deduces the class for the input data, completing the forward processing sequence as follows,
3.2. Extensible Codebook
While discretization mitigates data heterogeneity (see Ap- pendix A for theoretical analysis), it struggles with dis- tributions exhibiting significant variations. To effectively handle such heterogeneous data, we design an extensible codebook, beginning with a minimal set of codewords and progressively enlarging this set through a superior initial- ization strategy that benefits from our UEFL framework. This strategy facilitates stepwise mapping of diverse data distributions to distinct codewords. Starting with larger codebook can introduce uncertainty in codeword selection due to the concurrent training of multiple codewords.
Similar to VAE (Van Den Oord et al., 2017), we employ latent vectors as codewords, initializing a compact shared codebook with codewords c E R " where v represents the size of the initial codebook. The codewords are ini- tialized using a Gaussian distribution and shared across all data silos. After each iteration's 's uncertainty assessment, we determine which silos require additional codewords to im- prove prediction accuracy, and we extend the codebook ac- cordingly for these silos by adding v more codewords. The newly added codewords are initialized using K-means, lever- aging the encoder's improved latent features from the prior iteration to better align with the underlying data distribution. (2) To optimize codebook usage, data silos that demonstrated lower performance in the previous iteration are allowed to
2 = fox (z), c c= (e), p= (c)
where x, z, G, and p denote input data, latent features, dis- crete coded vectors, and the model prediction, separately.
After loss calculation, models undergo local updates through backpropagation. In manner akin to FedAvg (McMahan et al., 2017), these updated models are then relayed back to the server for a global update.
O2+ + 0 - - Vk:
K
T& 0k:
0 +-
(3) select codewords from both the newly added codewords and
72
k=
3

---

Uncertainty-Bant Beaetaintt--Bsedd Extensible-Codlcbook Federated Learning (UEFL)
ej > (1 + y)
Distributed Data
Discretization
Sample
Uncertainty Evaluator
1
Features
Codebook
$1
(0)
2
Low Uncertainty
(0)
Sample
Data
Shared Codebook
Oc
D2
$2
2
Encoder
High Uncertainty
Classifier
Data
Sample
Extended Codebooks
M
MM SMb
Initialization
Uncertain predictions
certain predictions
1st iteration
Next iterations
Compute loss C Concatenate
Figure 2: UEFL flowchart. In the first iteration, latent are mapped to initialized shared codewords by the discretizer OD- In the next iterations, UEFL identifies data from heterogeneous distributions with the uncertainty evaluator, and complements new codewords with K-means initialization to enhance the codebook. Clients with high uncertainty can select not only newly added codewords but also shared codewords.
the original shared codebook. Typically, the codebook only requires 1-3 extensions until all clients reach low uncertainty levels. The server updates the codebook by computing the average of codewords across the clients that utilize those specific codewords.
Segmented Codebooks. For complex datasets, a finite set of discrete codewords might not fully capture the diversity of image features. To bolster the robustness of our method- ology, we dissect features into smaller segments to pair them with multiple codewords, thus covering the entirety of fea- vector. This segmentation exponentially increases the codeword pool, ensuring robust representation capacity without necessitating large-scale increase and permitting efficient K-means- initialization. This design mini- mizes runtime overhead associated with larger codebooks.
For given iteration, if the codebook size for the kth client is Us, the feature vector is associated with codeword by the discretizer, which computes the distance between z and all available codewords, selecting the nearest one as follows,
i == arg min ||z -y cyll2
(4)
3.3. Loss Function
JE,.2...,
Since we introduce learnable codewords in our method, there are two parts of the loss function. For our task, we utilize cross-entropy as the loss function. For codebook op- timization, akin the strategy employed VQ- IQ-VAE (Van Den Oord et al., 2017), apply a stop gradient operation for the codeword update follows:
K-means Initialization. After the first iteration, the adapted encoder produces image features that more accurately re- flect the distribution of the training data. Instead of relying on random initialization methods like Gaussian distribution, we initialize new codewords using the centroids these fea- tures, obtained through K-means clustering. This approach expedites codebook training by providing more informed starting point for the new codewords, allowing them bet- ter align with the underlying data structure. As result, this initialization strategy facilitates faster convergence and improves the model's s ability to adapt to varying data dis- tributions across silos. This strategy hugely reduces the number training rounds required for model convergence (Details Appendix J).
Leode = =[>g(c) =1
Blle (gfl-=g(==1))) (3)
(5)
where z is the image latent features, c is discrete codewords, / is paree--arrmmttr to adjust the weights of two losses and sg(+) denotes the stop gradient function.
The total loss LV EFL the summation of Lrask and Leode-

---

Uncertainty-Bant -Based Extensible-Codlcbook Federated Learning (UEFL)
Algorithm I Uncertainty-Basee Extensible-CCdeeoook Fed-
effectively.
erated Learning (UEFL)
c=
Σ p plogp
(6)
Input: data distributions 1, D2.... Da
class
Parameters: uncertainty threshold y, learning rate 7, codewords loss weight 3 Sample K data silos from D1. D2. .... DM as clients
A low predictive entropy value signifies model confidence. whereas high value indicates increased uncertainty. For high entropy, introducing new codewords and conducting additional training rounds are essential steps. Given the vari- ability of uncertainty across datasets, establishing a fixed threshold is impractical. Instead, by analyzing all uncer- tainty values, we can benchmark against either the minimum or mean values pinpoint target silos. Our experiments showed superior results when using the minimum value as reference, thus guiding us to adopt the following threshold criterion:
Randomly initialize model parameters 0 and codebook with v codewords
Initially assign uncertainty for all clients be zero: ex
0, vk
repeat
for each round == 1, 2, .....d do
Broadcast 0 to all clients for K clients parallel do
if e > (1 + 7) (ennveee......(..... then K-means initialize another v codewords and add them codebook
C₁ S (1 (1++))
vjel.2,... (e)). Wk:
min
(7)
where y is a hyperparameter be set.
Update accesible codewords size for clients with high uncertainty: + +
Uncertainty decreases consistently during training, making it an effective stopping criterion for codebook extension. Be- sides, we manually set the maximum number of iterations to 5. A comparison with Deep Ensembles (Lakshminarayanan et al., 2017) is provided in Appendix E.
end if
Encode input into latent features: = S0g (x) Discretize latent features codewords Cs where i ar orgminjE1.2,... 2,...... -cyll2
Predict with coded vectors: p fee (c;) Compute codewords loss: Leods = ||sg(())
4. Experimental Results
e1 + B||c sg(2)
Compute output loss: Lrask - yogo
Experimental Setup. As discussed in (Kairouz al., 2021b; Zhou et al., 2023), there are two predominant forms of data heterogeneity in federated learning: feature het- erogeneity and label heterogeneity. Our UEFL focuses on tacking feature heterogeneity, and we mainly discuss fea- ture heterogeneity this section. The discussion for label heterogeneity and the comparison with VHL (Tang et al., 2022), FedBR (Guo al., 2023b) are in the Appendix I.
Update local parameters with gradient descent:
02 4 0g Vooodde + Ltask)
end for
Clients return all local models 02 to the server Update the server model 0 <- Σx= 1z end for
Evaluate uncertainty for each client with integrated model: 6 = Σ p pogg
Similar to Rotated MNIST (Ghifary et al., 2015), which cre- ates six domains through counter-clockwise rotations of 0", 15 30 45%, 60", and 75 on MNIST, we employ similar technique to introduce feature heterogeneity on five differ- ent datasets: MNIST, FMNIST, CIFAR1O, GTSRB, and CIFAR100, to validate our framework's robustness. In experiments, we create three domains by counter-clockwise rotating the datasets by 0' (D1), -50% (D2), and 120 (Dz). We sampled three data silos from each domain (i.e. totally 9 silos), and data silos for CIFAR 100 contain 4000 images each, while the other datasets consist of 2000 images per silo. Besides the regular training with multi-domain data silos, we also test UEFL for domain generalization (DG) task on Rotated MNIST (Ghifary al., 2015) and PACS (Li et al., 2017) datasets, which contains four distinct domains: art painting (A), cartoon (C), photo (P), and sketch (S).
Reduce the number of communication rounds
until eg < (1 + 7) minvje1,2....KK (ee)), Vk
3.4. Uncertainty Evaluation
As outlined in Section 3. evaluating model uncertainty is crucial for identifying data from heterogeneous distributions requiring supplementary codewords. In our work, we utilize Monte Carlo Dropout (MC Dropout) (Gal & Ghahramani, 2016) for uncertainty evaluation, incorporating two dropout layers into our model for regularization purposes. Unlike traditional usage where dropout layers are disabled during inference to stabilize predictions, we activate these layers during testing to generate variety of outcomes for uncer- tainty analysis. This variability is quantified using predictive entropy, as described in Equation (6), which serves to mea- sure the prediction dispersion across different evaluations
For RGB datasets like GTSRB, CIFARIO, 10, and CIFAR100, we adopt a pretrained VGG16 model in multi-domain train- ing. In contrast, for grayscale datasets such as MNIST and
5

---

Uncertainty-Baxxe Based Extensible-Codlcbook Federated Learning (UEFL)
The effect of discretization on FedAvg
Static Codebook v.s. Extensible Codebook
The effect of codebook initia lization
Accurasy
Uncertantyy
Par ciecity
1.0
Accurncy
Uncentinty
Perplexity
08-
0.8
Metric (a.u.)
Metric (a.u.
Metric (a.u.)
0.8
0.8
0.4
0.4
0.4
02
0.24
00
0.0
0.0
VO-FedAvg
(b)
(c)
Figure 3: Design of extensible codebook. (a) With the discretization (VQ-FedAvg), both accuracy and uncertainty get improved. (b) Our extensible codebook which starts from small capacity performs better than the static large codebook. (c) With K-means initialization, the utility of codewords (i.e. perplexity) gets significantly improved.
Table I: UEFL outperforms all baselines on heterogeneous data. DisTrans lacks Dropout layer, rendering it incapable of evaluating uncertainty. More results are in Appendix C.
MNIST
FMNIST GTSRB
CIFAR1O
CIFARI00
Methods
mA
mE
mA
mE
mA
mE
mA
mE
mA
mE
FedAvg (McMahan et al., 2017)
0.780
0.273
0.803
0.273
0.660
0.636
0.617
0.177
0,088
1.91
DisTrans (Yuan et al., 2022)
0.815
0.707
0.898
0.699
0.267
UEFL (Ours)
0.920
0.149
0.850
0.167
0.942
0.0239
0.720
0.0222
0.326
0.655
FMNIST, lacking pretrained models, we design a convolu- tional network comprising three ResNet blocks, training it from scratch. And for DG, adopt pretrained ResNet18 for both datasets. Initial codebook sizes are set to 32 for MNIST and 64 for the remaining datasets, with an equiv- alent number of codewords added in each subsequent iter- ation. While additional iterations may converge within 5 rounds, we extend this to 20 enhanced experimental clar- ity. The uncertainty evaluation conducted 20 times using dropout rate of 0.1, with thresholds y 0.3 for MNIST, 0.1 for FMNIST, GTSRB, and CIFAR 100, and 0.2 for CI- FAR10, to fine-tune performance. These experiments are performed on machine with two NVIDIA A6000 GPUs.
demonstrating the effectiveness of feature discretization in enhancing performance within heterogeneous federated learning context, shown in Figure ((a).
Extensible Codebook v.s. Static Large Codebook. To vali- date our extensible codebook's superiority over starting with large codebook, we ensured both methods ended with the same number of codewords through experiments. Results on CIFAR 100 in Figure 3(b) demonstrate the difficulties associated with a larger initial codebook in codeword selee- tion for image features. Conversely, gradually expanding the codebook significantly improved codeword differentia- tion, yielding better outcomes, such as enhanced accuracy (from 0.13 to 0.34), reduced uncertainty (0.78 vs. 1.66 for the static approach), and increased utilization codewords.
Evaluation Metrics. We compute the mean Top-1 accuracy (mA), mean entropy (mE), and mean perplexity (mP) across all clients facilitate a direct comparison, where perplexity quantifies the utility of codewords.
Codebook Initialization. Section 3.1 highlights our UEFL's capability for efficient codeword initialization via K-means, utilizing features from a finetuned encoder. The efficacy of initialization is validated in Figure 3(c) with re- sults on MNIST, showing enhancements across all metrics.
4.1. Extensible Codebook
Discretization for Heterogeneous FL. To show the effect- of discretization to tackle the data heterogeneity in FL, we design toy experiment MNIST. Temporarily setting aside federated learning's privacy considerations, we directly discretized the features for each client using the dis- tinct codebooks based on its originating domain. With this discretization VQ-FedAvg, mean accuracy was im- proved from 0.834 0.907 with the reduction of uncertainty,
4.2. UEFL for Multi-Domain Learning
We conducted comparative experiments on five datasets with introduced feature heterogeneity against leading algorithms, specifically the baseline FedAvg (McMahan et al., 2017) and DisTrans (Yuan et al., 2022). For accuracy compari- son, DisTrans generally exhibits better performance than FedAvg, making our primary point of comparison. Re-
6

---

Uneertainty-Baase Extensible-Codlcbook Federated Learning (UEFL)
Anturacy tifferent dalla silos
Uncerisiinly on enl dbala sillos
Per Perplexily on differeni datas siloes
Fosrin
Foddig
0.3
100
Uncertaint
Paralouuty
66
0.2-
0.14
16
92
32
1c
25
680
sig
(a) Accuracy comparison.
(b) Uncertainty comparison.
(c) Perplexity comparison.
Figure 4: Experiments
MNIST. D3 presents much lower accuracy and higher uncertainty compared D1, D2 for FedAvg. And the perplexity results show that our UEFL assigns new codewords to D3 to improve the performance.
Table 2: Comparison with different methods for DG. Results are on six domains of Rotated MNIST, four domains of PACS and their average. Our approach is compared with baselines: FedAvg, FedSR, and FedIIR.
Rotated MNIST
PACS
Methods
Mo
M15
Ma0
Mas
M60
M75
Ave.
Ave.
FedAvg (McMahan et al., 2017)
82.7
98.2
99
99.1
98.2
89.9
94.5
78
73
92
79
80.3
FedSR (Nguyen et al., 2022)
84.2
98.0
98.9
99.0
98.3
90.0
94.7
83
75
94
82
83.4
FedlIR (Guo et al., 2023a)
83.8
98.2
99.1
99.1
98.5
90.8
95.0
83
76
94
82
83.7
UEFL (ours)
88.1
97.3
97.6
97.8
97.9
93.2
95.3
81
80
94
82
84.5
garding uncertainty comparison, because DisTrans lacks Dropout layers, precluding uncertainty evaluation, clusively compare uncertainty metrics with FedAvg.
4.3. UEFL for Domain Generalization (DG)
For DG task, the trained model needs be evaluated on out-of-distribution domain and we follow the evaluation method in (Nguyen et al., 2022; Guo et al., 2023a). Specif- ically, we perform "leave-one-domain-oo where we choose one domain as the target domain, train the model on all remaining domains, and evaluate it on the chosen domain. Each source domain is treated as a client.
The results in Table I provide a comprehen- sive comparison, illustrating that UEFL. surpasses all other SOTA methods in both accuracy and uncertainty reduction. Specifically, UEFL achieves accuracy improvements rang- ing from 3% 22.1% over DisTrans. Our UEFL improves uncertainty compared to FedAvg, achieving reductions 38.83%-96.22 2.%... Figures 4(a) and 4(b) details performance across each data silo. highlighting UEFL's effectiveness in elevating the accuracy and degrading the uncertainty.
Performance.
experiments,
As shown in Table 2, our UEFL enhanced mean accuracy the RotatedMNIST dataset, elevating it from 0.945 to 0.953. This performance exceeds that of FedSR (Nguyen et al., 2022) at 0.947 and FedlIR (Guo et al., 2023a) at 0.95, Similarly, on the PACS dataset, UEFL improved mean accu- racy from 0.803 to 0.8453, surpassing FedSR's s 0.834 and FedIR 0.837. These results underscore UEFL's efficacy in tackling feature heterogeneity and superior performance the federated DG task, beating state-of-the-arr methods.
Codewords Perplexity. Figure 4(c) presents perplexity comparison between our UEFL and FedAvg, illustrating enhanced codebook utilization after assigning new code- words to D3. This adjustment not only benefits Dy but also improves the codebook utilization for D, and 'D2.
Computation Overhead. Our approach introduces only a small codebook, thus incurring negligible memory and computational overheads. Specifically, for the CIFAR dataset, the parameter count for the baseline FedAvg model 14.991M, whereas our UEFL model slightly increases to 15.491M indicating tiny memory increment of 3.34%. In terms of runtime, UEFL also exhibits minimal increase from 16. 154ms to 16.733ms (3.58% increase). These find- ings underscore UEFL'S suitability for deployment on edge devices. More details are included in Appendix G.
4.4. Ablation Study
Imbalanced Clients. We constructed an experimental setup with three data silos from D, and one each from D2 and Dz. totaling five silos. Our UEFL can also improve both accu- racy (from 0.508 to 0.828) and uncertainty (from 0.256 to 0. 105) in this scenario. Detailed results are in Appendix K.
Large Number of Clients. We follow (Guo et al., 2023a) to further partition the five training domains Rotated MNIST into 50 sub-domaing and adopt the setup from (Zhang et al.,

---

Secertainty-Basee Extensible-Codccbok Federated Learning (UEFL)
Table 3: UEFL is scalable for 50 clients and beats all SOTA methods: FedAvg, FedSR and FedIIR.
Domains
Methods
#Clients
Backbone
Average
Mo
M15
Mao
M45
Moa
M75
FedAvg (McMahan al., 2017)
50
ResNet18
77.9
95.9
96.9
97
96
81.2
90.8
FedSR (Nguyen et al., 2022)
50
ResNet18
78.3
95.7
96.3
97.1
96
84
91.2
FedlIR (Guo et al., 2023a)
50
ResNet18
84
96.8
97.7
97.7
97.4
84.5
93
UEFL (ours)
50
ResNet18
86.4
95.5
96.4
96.9
94.7
90.6
93.42
Table 4: UEFL also scalable for 100 clients and beats all SOTA methods following the experimental setup in FedCR.
Methods
#Clients
EMNIST-L
FMNIST
CIFARI00 CIFARIO
FedAvg (McMahan et al., 2017)
100
95.89
88.15
76.83
32.08
FedSR (Nguyen et al., 2022)
100
86.22
85.55
61.47
40.82
FedCR (Zhang et al., 2023)
100
97.47
93.78
84.74
62.96
UEFL (ours)
100
98.29
93.93
86.11
63.37
Performance difterent codewords
Performance dilferent segments
Performance different thresholds
1.0-
1.0
Metric (a.u.)
0.8-
Metric (a.u.)
Accuracy 3PARTO: Aotra Ureedainn/ Uneembny
Metric (a.u.)
0.8-
0.8
0.8-
0.4-
Acsurasy Unnemwinty Parplealt
0.4
Accurity aoy
Unnerrsiryy Perplestt
0.2
0.-
0.0
32
256
0.1
0.3
0.4
0.5
0.8
Codebook Size
Number of Segments
Threshold
(b)
Figure 5: (a) 32 codewords are sufficient. (b) CIFAR100 requires segments. (c) A smaller threshold performs better.
2023) to evaluate performance on 100 clients. both cases, UEFL achieves the highest mean accuracy, outperforming all baselines, shown in Table 3 and Table 4, demonstrating its scalability larger number of clients.
increase computational overhead. Thus, in such cases, there is trade-off between runtime and performance.
5. Conclusion
Number of codewords and segments. We investigate the impact of varying the number of initialized codewords in our extensible codebook, to balance accuracy with runtime effi- ciency K-means initialization. In Figure 5(a) for GTSRB, initializing with 32 codewords provides comparable accu- racy and uncertainty metrics. For more complex datasets, we enhance selection capacity using codeword segmenta- tion. Figure 5(b) demonstrates that segmenting codewords into 4 parts leads to enhanced performance on CIFAR
In this work, we address the challenge of data heterogeneity among silos within federated learning setting by introducing an innovative solution: an extensible codebook designed to map distinct data distributions using varied codeword pools. Our proposed framework, Uncertainty-Base Extensible- Codebook Federated Learning (UEFL), leverages this ex- tensible codebook through an iterative process that adeptly identifies data from unknown distributions via uncertainty evaluation and enriches the codebook with newly initial- ized codewords tailored to these distributions. The iterative nature of UEFL, coupled with efficient codeword initializa- tion using K-means, ensures codewords closely matched with the actual data distribution, thereby expediting model convergence. This approach allows UEFL to rapidly adjust to new and unseen data distributions, enhancing adaptabil- ity. Our comprehensive evaluation across various prominent datasets showcases UEFL's effectiveness.
Uncertainty Threshold. our UEFL, the uncertainty evaluator plays pivotal role in identifying heterogeneous data without needing direct data access, with the threshold selection being critical. As illustrated in Figure 5(c), lower threshold imposes stricter criteria, pushing the model to achieve higher performance. However, it's important to recognize that beyond a certain point, further reducing the threshold may not significantly enhance outcomes but will
8

---

Uncertainty- Extensible-Codccbok Federated Learning (UEFL)
Acknowledgements
Ghifary, M., Kleijn, W. B., Zhang, M., and Balduzzi, D. Do- main generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE international conference on computer vision, pp. 2551-2559, 2015.
This work was supported in part by the joint funding pro- gram between the National University of Singapore and the University of Toronto. We gratefully acknowledge their support.
Ghosh, A., Chung. J., Yin, D., and Ramchandran, K. An efficient framework for clustered federated learning. Ad- vances in Neural Information Processing Systems, 33: 19586-19597 2020.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Guo, Y., Guo, K., Cao, X., Wu, T., and Chang, Out- of-distribution generalization of federated learning via implicit invariant relationships. In International Confer- ence on Machine Learning, pp. 11905-11933, PMLR, 2023a.
References
Guo, Y., Tang, X., and Lin, T. Fedbr: Improving feder- ated learning on heterogeneous data via local learning bias reduction. In International Conference on Machine Learning, pp. 22034-12054 PMLR, 2023b.
Agarwal, N., Kairouz, P., and Liu, Z. The skellam mech- anism for differentially private federated learning. Ad- vances in Neural Informatio Processing Systems, 34: 5052-5064, 2021.
Hard, A., Rao, K., Mathews, R., Ramaswamy, S., Beaufays, F., Augenstein, S., Eichner, H., Kiddon, C., and Ramage, D. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: /8/ 1.06604, 2018.
Ahn, H., Cha. S., Lee, D., and Moon, T. Uncertainty- based continual learning with adaptive regularization. Advances in neural information processing systems, 32, 2019.
Blundell, C., Comebise,, J., Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Huang, W. Ye. M., and Du B. Learn from others and be yourself in heterogeneous federated learning. In Proceed- ings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10143-10153, 2022.
Chen, F. Luo, M., Dong, Z., Li, Z., and He, X. Feder- ated meta-learning with fast convergence and efficient communication. arXiv preprint DeXivv/802.07000
Kairouz, P., Liu, Z., and Steinke, T. The distributed discrete gaussian mechanism for federated learning with secure aggregation. In International Conference Machine Learning, pp. 5201-5212 PMLR, 2021a.
2018.
Chen, Q., Whitbrook, A., Aickelin, U., and Roadknight, C. Data classification using the dempster-shafe method. Journal of Experimental Theoretical Artificial Intelli- gence, 26(4):493-517, 2014.
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, 7., Cormode, G., Cummings, R., et al. Advances and open problems in federated learning. Foundations and trends@ in machine learning. 14(1-2):1-210, 2021b.
Dusenberry, M. W., Tran, D., Choi, E., Kemp, J., Nixon, J., Jerfel, G., Heller, K., and Dai, A. M. Analyzing the role of model uncertainty for electronic health records. In Proceedings of the ACM Conference Health, Inference, and Learning, pp. 204-213 2020.
Kalra, S., Wen, J.. Cresswell, J. C., Volkovs, M., and Tizhoosh, H. Decentralized federated learning through proxy model sharing. Nature communications, 14(1): 2899, 2023.
Gal, Y. and Ghahramani, Z. Dropout as a bayesian approx- imation: Representing model uncertainty in deep learn- ing. In international conference on machine learning, pp. 1050-1059. PMLR, 2016.
Kendall, A. and Gal, Y. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Gawlikowski, J., Tassi, C. N., Ali, M., Lee, J., Humt, M., Feng, J., Kruspe, A., Triebel, R., Jung, P., Roscher, R., et al. A survey of uncertainty in deep neural networks. 2021.
Konečný, J., McMahan, H. B., Ramage, D., and Richtárik, P. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint 2016.
arXiv preprint arXiv:2107.0334
Geyer, R. C., Klein, T., and Nabi, M. Differentially private federated learning: A client level perspective. arXiv preprint
Lahlou. S., Jain, M., Nekoci, H., Butoi, V. I., Bertin, P., Rector- Rector-Brooks, J., Korablyov, M., and Bengio, Deup:
2017.

---

Encertainty-Basel Extensible-Codcbook Federated Learning (UEFL)
Direct epistemic uncertainty prediction. arXiv preprinl arXiv:2102.08501, 2021.
Tang, Z.,, Zhang, Y., Shi, S., He, X., Han, B., and Chu, X. Virtual homogeneity learning: Defending against data heterogeneity in federated learning. In International Con- ference on Machine Learning, pp. 21111-1111111 PMLR, 2022.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Yang, Q., Liu, Y., Chen, T., and Tong, Y. Federated machine learning: Concept and applications. ACM Transactions Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
Li, D., Yang, Y., Song, Y.-Z., and Hospedales, T. M. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision, pp. 5542-5550, 2017.
Yuan, H., Hui, B., Yang, Y., Burlina, P., Gong, N. Z., and Cao, Y. Addressing heterogeneity in federated learning via distributional transformation. In European Confer- ence Computer Vision, pp. 179-195. Springer, 2022.
Li, L., Xu, W., Chen, T., Giannakis, G. B., and Ling, Q- Rsa: Byzantine-robust stochastic aggregation methods for distributed learning from heterogeneous datasets. In Proceedings of the AAA/ Conference on Artificial Intelli- gence, volume 33, pp. 1544-1551, 2019.
Zhang, H., Li, C., Dai, W., Zou, J., and Xiong, H. Feder: Personalized federated learning based on across-client common representation with conditional mutual infor- mation regularization. In International Conference on Machine Learning, pp. 41314-41330, PMLR, 2023.
Li, T., Sahu, A. K., Sanjabi, M., Zaheer, M., Talwalkar, A., and Smith, V. On the convergence of federated op- timization in heterogeneous networks. arXiv preprint orXiv:1812.0611 2018.
Zhang, T., Zhang, S., Chen, Z., Bengio, Y., and Liu, D. Pmfl: Partial meta-federated learning for heterogeneous tasks and its applications on real-world medical records. In 2022 IEEE International Conference on Big Data (Big Data), pp.
Liu, D., Lamb, A., Kawaguchi. K., Goyal, A., Sun, C.,, Mozer, M., and Bengio, Y. Discrete-valuee neural com- munication in structured architectures enhances general- ization. 2021.
IEEE, 2022.
Zhao, Y., Li, M., Lai, L., Suda, N., Civin, D., and Chandra, V. Federated learning with non-iid data, arXiv preprint arXiv: 11880000022 2018
Louizos, C. and Welling, M. Multiplicative normalizing flows for variational bayesian neural networks. In Interna- tional Conference on Machine Learning, pp. 2218-2221 PMLR, 2017.
Zhou, T., Zhang, J., and Tsang, D. H. Fedfa: Federated learning with feature anchors align features and clas- sifiers for heterogeneous data. IEEE Transactions Mobile Computing, 2023.
McMahan, B., Moore, E., Ramage, D., Hampson, S., and
Arcas, B. A. Communication-efficeee learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Nado, Z., Band, N., Collier, M., Djolonga, J., Dusenberry, M. W., Farquhar, S., Feng, Q. Filos, A., Havasi, M., Jenatton, R., al. Uncertainty baselines: Benchmarks for uncertainty & robustness in deep learning. arXiv preprint arXiv:2106.0400 04015, 2021.
Nguyen, A. T., Torr, P., and Lim, S. N. Fedsr: A simple and effective domain generalization method for federated learning. Advances in Neural Information Processing Systems, 35:38831-38841 2022.
Papyan, V., Han. X., and Donoho, D. L. Prevalence of neural collapse during the terminal phase of deep learn- ing training. Proceedings of the National Academy of Sciences, 100,,111111 117(40):2452-24666) 2020.
10

---

Uncertainty-Bast -Based Extensible-Codlcbokk Federated Learning (UEFL)
A. Discretization for Generalization
Theoretically, the inclusion of the discretization process offers two key advantages: (I) enhanced noise robustness, and (2)
reduced underlying dimensionality. These benefits are demonstrated in the following two theorems. (Liu et al., 2021).
Notation: h is input vector, h E H Rm. L is the size of codebook, G' is the number of segments, q(-) is discretization process, (f-) is any function (model), Given any family of sets S = (S,, Sk) } with Si, Sx S H. we define of by of K E S& (hh) for all & [K], where [K] {1.... K). And we denote by (Qk)kE[LO) all the codewords.
Theorem 1: (with discretization) Let Sk {Qx} for all & E [L.. Then, for any ó O, with probability at least 1 5 over on iid draw of n eramples (h.) the following holds for any &b R + R. and all k E [7. if j (h)l S a for all h E H. then
GIn(L) + In(2/6)
EA lo (g(h, L, L,)
Σ 1=1
o (a L, G)) = Oop
(8)
n
2n
where no constant is hidden in O.
Theorem 2: (without discretization) Assume that ||h||2 I Ru for all h H RT. Fix C e argminz
C RTh H S UeezBlel) where B(c] = (ær E Rm : | cll2 s Rx/(2/T)). Let Sp = Blea] for all k E [[C]] where Cs E and Us(cu} = C. Then, for any ô > 0, with probability at least - over an iid draw of
eramples (h,)"=,,the following holds for any Φ Rm + R and all & E [[]]] if 1 (h) s a for all H and
| (h) - o% (h') & sa hh - h'll2 for all h h' E Sk. for all h, h' E Sk, then
n1
mln(4 ((44//m)))))) + In(2/8)
Sh SRRH Vn
EA LO (h))
(h:)
O(aa
(9)
2n
where no constant is hidden in and 5e = Su + Σ. (h; E B(c..)).
Based on these two theorems, we can determine that the performance gap between training and test data is smaller when discretization is applied, due the following two points:
There is an additional error without discretization (i.e. St vR
in the bound of Theorem This error disappears with discretization in the bound of Theorem as the discretization process reduces the sensitivity to noise.
The discretization process reduces the underlying dimensionality of mln(44/mm) without discretization (in Theorem 2) to that of GIn(L) with discretization (in Theorem 1). Since the number of discretization heads G (eg. G 1, or in our case) always much smaller than the number of dimensions m, the inequality GIn(L) S mln(44/mm) consistently holds.
B. Discretization for non-iid Federated Learning
Notation:
Consider a federated setting with K clients. Client & has ny samples independently drawn from its own distribution Ps. Total samples of all clients n = Σk=1 ni. The global distribution -K na Pa.
Theorem 3: (with discretization) if |0g (h) S for all h E H. then, for any 8 0, with probability at least
K n
Gln I. (n(2//8) (ln22k1((//)))
nk En (gh L, n
ΣΣ
lggh
(K) L.G)) = =Oaay
(10)
n
2n
vn
where v6 = 1 Div ) denoting the KL divergence between the client distributions and the global distribution after discretization.
K

---

Eneertainty-Base
Extensible-Codlcbook Federated Learning (UEFL)
Theorem 4: (without discretization) Assume that |
Ryg then for any with probability 2
Tk Σ
En P2 (6x
(h Ooay
(18)
mln(4 + In(2K/6)
sRH + U)
(11)
n
2n
vn
where V E Div(P P), demoting the KL divergence between the client distributions and the global distribution.
The KL divergence term p(4) after discretization significantly smaller than the original // due discretization-induds invariance. Therefore, discretization not only improves robustness noise and reduces dimensionality, but importantly, it explicitly and effectively mitigates the adverse effects of data heterogeneity typical in non-IID federated learning environments.
C. Different data heterogeneity for UEFL
C.1. UEFL for Multi-Domain Learning
Table 5: UEFL outperforms all baselines on heterogeneous data. DisTrans lacks a Dropout layer, rendering incapable of evaluating uncertainty. CIFAR100" exhibits poor performance due to the highly heterogeneous experimental setup.
MNIST
FMNIST
GTSRB
CIFARIO
CIFAR100
Methods
Data
mA
mE
mA
mE
mA
mE
mA
mE
mA
mE
D 0.874
0.212
0.801
0.246
0.670
0.623
0.676
0.172
0.110
1.74
D2
0.848
0.231
0.825
0.232
0.677
0.634
0.622
0.178
0.072
1.86
FedAvg
D3
0.618
0.377
0.784
0.341
0.634
0.652
0.553
0.183
0.083
2.13
All
0.780
0.273
0.803
0.273
0.660
0.636
0.617
0.177
0.088
1.91
D1
0.856
0.721
0.898
0.721
0.289
D2
0.799
0.705
0.900
0.719
0.261
DisTrans
D3
0.789
0.694
0.897
0.659
0.251
All
0.815
0.707
0.898
0.699
0.267
D1
0.951
0.120
0.857
0.147
0.95
0.0196
0.776
0.0192
0.362
0.728
D2
0.885
0.196
0.848
0.188
0.964
0.0206
0.713
0.0245
0.335
0.624
UEFL (Ours)
D;
0.924
0.131
0.845
0.167
0.911
0.0314
0.671
0.0229
0.282
0.612
All
0.920
0.149
0.850
0.167
0.942
0.0239
0.720
0.0222
0.326
0.655
C.2. Lower Data heterogeneity on CIFAR100
Table 5, CIFAR100 exhibits poor performance due highly heterogeneous experimental setup. provide a detailed comparison with the baseline, we used VGG16 backbone to test CIFARI00 under multiple settings: (1) Local Training: all data trained together without a distributed setting; (2) FedAvg (w/o hete): CIFAR 100 split into 5 clients, each with 10,000 images, to evaluate FedAvg performance; (3) FedAvg (w/ hete): images for the clients rotated by -309, -15%, 0", 15%. and 30", respectively, to introduce data heterogeneity, and FedAvg performance was evaluated; (4) UEFL (w/ hete): tested under the same heterogencous setup. We trained models from scratch and with pre-trained weights. Results are presented in Table 6, showing that UEFL consistently outperforms FedAvg in both cases.
C.3. Different data heterogeneity on CIFAR100
By progressively increasing the rotation angles to simulate greater data heterogeneity, we evaluated UEFL's performance under varying levels heterogeneity. As shown in Table 7, while overall performance decreases with higher heterogeneity, UEFL consistently outperforms FedAvg, with the performance gap widening as heterogeneity increases, demonstrating UEFL's superiority in addressing data heterogeneity.
12

---

Uncertainty-Bant Based Extensible-Codccbokk Federated Learning (UEFL)
Table 6: Under lower data heterogeneity on CIFAR 100, UEFL continues to outperform FedAvg for both training from scratch and using pre-trained weights.
Training Strategy
Local training
FedAvg (w/ hete)
FedAvg (w/ hete)
UEFL hete)
From scratch
0.3852
0.2447
0.0852 0.5005
0.1062
Pre-trained
0.6604
0.6496
0.5619
Table 7: With different rotation angles, UEFL keep outperforms FedAvg.
Rotation Angles
FedAvg
UEFL
f-10%, -5, 5%, 10%}
0,5935
0.6232
(-20%,-10%, 0% 10", 20")
0,5469
0.6039
(-30%.-15*, 15%, 30")
0,5106
0,56
{-40,, 20, 40%)
0,4799
0.5311
(-50%. 25, 25% 50%)
0,4494
0.5074
--60 00, 60%)
0.4368
0.4939
0",
0.4188
0.4737
D. Optimal training epochs of baselines
To fully demonstrate the efficacy of our UEFL, besides evaluating baselines with the same total training epochs UEFL, remove the additional training epochs from UEFL iterations and obtain the optimal performance, for more fair comparison.
Table 8: UEFL outperforms all baselines without additional training epochs.
MNIST
FMNIST
GTSRB
CIFAR10
CIFARI00
Methods
mA
mE
mA
mE
mA
mE
mA
mE
mA
mE
FedAvg
0.782 0.920
0.261
0.801
0.289
0.657
0.645
0.618
0.173
0.093
1.74
UEFL (Ours)
0.149
0.850
0.167
0.942
0.0239
0.720
0.0222
0.326
0.655
E. Deep Ensembles
We evaluated Deep Ensembles (Lakshminarayanan et al., 2017) by creating 5 ensembles to assess uncertainty for our method. Table shows a comparison between Deep Ensembles and Monte Carlo Dropout. From the results, we observe that the accuracy when using Deep Ensembles is quite similar Monte Carlo Dropout, apart from the stochastic variations. This is expected, the accuracy not directly impacted by the choice of uncertainty evaluation method. However, the uncertainty values for Deep Ensembles are higher than those for Monte Carlo Dropout, likely due to the use of only 5 ensembles for evaluation to reduce computational time. In conclusion, while the accuracy is comparable, Deep Ensembles require significantly more computational resources due to the need to train multiple networks. Therefore, Monte Carlo Dropout is a more efficient and suitable choice for our approach.
F. Neural Collapse
We conducted experiments MNIST and CIFAR-100 based the framework presented in (Papyan et al., 2020). According to the paper, when training continues until the training error reaches 0 (i.e. training accuracy exceeds 99.99. for MNIIT/CIFARR-O)) the Terminal Phase of Training (TPT) begins, during which neural collapse (NC) emerges. To validate this, we first conducted local training by training on all data together. Our results confirmed the paper's 's claim that additional training beyond the zero-error point leads to improved performance. We then extended these experiments to the federated learning setting. The detailed results are presented in Table 10.
13

---

Snerttainty-Basee Extensible-Codccbokk Federated Learning (UEFL)
Table 9: Comparison of Deep Ensembles and Monte Carlo Dropout for uncertainty evaluation.
MNIST
FMNIST
GTSRB
CIFAR1O
CIFAR100
Methods
mA
mE
mA
mE
mA
mE
mA
mE
mA
mE
Monte Carlo Dropout
0.920
0.149
0.850
0.167
0.942
0.0239
0.720
0.0222
0.326
0.655
Deep Ensemble
0.926
0.211
0.853
0.289
0.940
0.041
0.717
0.043
0.331
0.873
From these experiments, we observed the following findings: 1. The dropout layer must be removed; otherwise, the training accuracy cannot exceed 99.9% (e.g., the final training accuracy for MNIST is limited to 96% with dropout). 2. Compared to local training, more training epochs are required to reach the TPT in federated learning. For example, on CIFAR-100 with data heterogeneity, federated learning requires rounds of training (44 x epochs), while local training achieves TPT in 38 epochs. 3. Although neural collapse yields improved performance, UEFL consistently outperforms especially on more complex datasets like CIFAR-100. This is partly because removing dropout layers for neural collapse increases the risk of overfitting.
Table 10: Comparison with neural collapse (NC).
Method
MNIST
CIFAR
zero-error
0.9375
0.3473
last epoch (NC)
0.9584
0.3482
UEFL (Ours)
0.9778
0.5074
G. Computation Cost
Table 1l compares FedAvg and UEFL. The parameter count for UEFL in this table represents the final model size, including 256 codewords after two iterations starting from 64. In our experiments, 256 is the largest final codebook size across all datasets. For simpler datasets like MNIST, UEFL demonstrates an even lower computation cost.
Table 11: With different rotation angles, UEFL keep outperforms FedAvg.
Method
#Params (M)
CPU runtime (ms)
GPU runtime (ms)
FedAvg
14.991 15,491
16.102
16.154
UEFL (Ours)
16.611
16.733
H. Convergence Curves
For the experiments on the CIFARIO dataset, in Figure 6, at round 40, after we assign new codewords, rapid performance gains are evident. Remarkably, the training process demonstrates swift convergence, typically within just five rounds. For illustrative clarity and to underscore the differential impact, we extend the training to 20 rounds in subsequent iterations, showcasing the accelerated and effective adaptation of our approach. In addition, after 60 rounds, even if we keep adding new codewords, the increased perplexity denotes a higher utilization of the codebook. However, there is no significant improvement in accuracy or uncertainty. Figure also shows large boost with our UEFL on MNIST.
For the experiments on the FMNIST and GTSRB datasets, as depicted in Figure 7, we introduce new codewords to the codebook only once. Notably, there is a clear "performance jump" evident in all six figures, showcasing the rapid adaptation of our UEFL to new data distributions.
14

---

Uncertainty-Bant -asted Extensible-Codccbok Federated Learning (UEFL)
(6)
(b)
[c]
(d)
Figure 6: Learning curves of MNIST and CIFARIO. (a) Accuracy of MNIST. (b) Accuracy of CIFARIO. (c) Uncertainty of CIFAR10. (d) Perplexity of CIFARIO.
Accuracy
Uncertainty
Codeword Loss
05011
45311
325
89:35
AC0100
01-18
36-
030
04-
0.2
011
23
28-
7.-
03-
Figure 7: Learning curves of FMNIST and GTSRB. (left) Accuracy results. (mid) Uncertainty results. (right) codeword loss results.
I. Our UEFL for Label Heterogeneity
Similar to (Tang et al., 2022; Guo al., 2023b), we introduce label heterogeneity with dirichlet distribution (a 0.1). The results in Table 12 show that our UEFL can also tackle the label heterogeneity when compared to FedAvg and performs better than VHL for CIFARIO even if it cannot perform well as FedBR.
Table 12: Comparison with different methods on data with label heterogeneity.
FMNIST
CIFAR10
Method
FedAvg
VHL
UEFL (ours)
FedAvg
VHL
FedBR
UEFL (ours)
mA
87.45
91.52
90.59
58.99
61.23
64.61
62.67
J. K-means Initialization
Figure illustrates this concept: gray points represent features from the trained encoder, clustered according to their data distributions. While direct data access is restricted, differentiation by uncertainty allows us to identify and utilize the centroids of these clusters via K-means for codeword initialization.
15

---

Uinertaintt-Baace Extensible-Codecbook Federated Learning (UEFL)
Table 13: Our UEFL improves the performance of unbalanced data.
FedAvg
UEFL
Data
Silo
Acc
Entropy
Acc
Entropy
aaa
0.964
0.0312
0.952
0.0291
D
S1.
0,936
0,0252
0.974
0,0170
$1,
0.964
0.0499
0.944
0.0308
D2
$2.
0.796
0.1477
0.836
0.1261
D3
$3a
0.508
0.2560
0.828
0.1048
And to bolster the robustness of our methodology, we dissect features into smaller segments uugments-using factors like 2 or 4-to pair them with multiple codewords, thus covering the entirety of feature vector as illustrated in Figure 3. This segmentation exponentially increases the codeword pool n² or n², ensuring robust representation capacity.
Ci
2
- cjll2
Zip
Cip
||
Segment
Concatenate
argmin
je(1.2.....7)
Zia
Cia
Feature segments
Codewords for D
Codewords for 2
Codewords for 2D3
(a) Kmeans initialization
(b) Discretizer
Figure 8: (a) Kmeans initialization for heterogeneous data silos. (b) Workflow of discretizer.
K. Imbalanced Clients
As shown in Table 13, our UEFL also works for imbalanced data silos when there three clients sampled from the same domain. Both accuracy and uncertainty get improved, especially for the third domain.
L. UEFL Optimization
Number of Codewords. We investigate the impact varying the number initialized codewords within our extensible codebook Table 14, aiming strike a balance between achieving competitive accuracy and optimizing the runtime efficiency of the K-means initialization. Our findings, for the GTSRB dataset, reveal that starting with 32 64 codewords offers comparable accuracy and uncertainty metrics to larger codebooks, while significantly enhancing the efficiency of the K-means initialization. This efficiency highlights the efficacy of our proposed approach.
In addition, for more complex datasets, requiring a broader representation of image features but with minimal initialization time, we employ codeword segmentation to enhance selection capacity efficiently. We explore the impact of segmentation factors of 1, 2, and 4, starting with 16 codewords GTSRB and 32 for CIFAR100. Our findings indicate that, particularly
16

---

JincertinttyBaaxe Extensible-Codlcbook Federated Learning (UEFL)
for CIFAR100, splitting vectors into 4 segments with only 32 initialized codewords achieves impressive performance. Similarly, for GTSRB, segmentation into 2 parts is adequate for effective image feature representation.
#Codes
Data
code +
mPt
mE,
mA
codebook growth
D
3.61
4.78
2.02
0.257
8
D2
3.68
4.86
2.01
0.265
16 32 64
D3
3.38
4.81
2.03
0.249
D1
3.41
10.46
1.36
0.515
16
D2
3.55
10.54
1.37
0.506
16 32 64
D3
3.23
10.61
1.39
0.486
D1
0.127
25.91
0.0412
0.956
32
D2
0.0885
25.42
0.0313
0.966
32 - 64 + 128
D3
0.178
26.37
0.117
0.911
D1
0.0975
26.79
0.0086
0.965
64
D2
0.0853
26.32
0.0084
0.974
64 128
D3
0.1907
27.23
0.0166
0.926
D1 0.0512
38.73
0.954
128
D2
0.0453
34.16
0.968
128 256
D3
0.0726
43.42
0,917
Di
0.0543
41.20
0.0043
0.962
256
D2
0.0301
38.96
0.0054
0.959
256 + 512
DD
0.0577
50.57
0.0103
0.904
Table 14: Number of codewords. Experiment are on GTSRB dataset. "." denotes value close to
Codebook Initialization. Section 3.1 highlights UEFL framework's capability for efficient codeword initialization via K-means, utilizing features from a trained encoder. The efficacy of K-means initialization is validated in Table 15 with results from the MNIST dataset, showing enhancements across all metrics.
Extensible Codebook v.s. Static Large Codebook. To validate our extensible codebook's s superiority over starting with a large codebook, we ensured both methods ended with the same number of codewords through experiments. For the CIFAR100 dataset, the extensible codebook was initially set 128 codewords and expanded twice, while the static codebook was fixed at 512 codewords. Results showcased in Table 16 demonstrate the difficulties associated with larger initial codebook in codeword selection for image features. Conversely, gradually expanding the codebook significantly improved codeword differentiation, yielding better outcomes, such enhanced accuracy (0.375 for Domain 1) and reduced uncertainty (0.78 vs. 1.66 for the static approach). In addition, perplexity results reveal increased utilization of our extensible codebook, offering clear evidence of our design's 's superiority.
Different Uncertainty Threshold. our UEFL,, the uncertainty evaluator plays pivotal role identifying heterogeneous data without needing direct data access, with the threshold selection being critical. An optimal threshold enhances the model's ability distinguish between data silos, leading quicker convergence. As illustrated in Figure 9, lower threshold
Table 15:
The experiments were conducted on the MNIST dataset with 128 initialized codewords and segmentation factor 1.
The model with K-means initialization outperforms without
Codebook
Data
Leode
mP
mE
mA
D1
0.6202
6.26
0.125
0.888
w/o init
D2
0.6063
5.99
0.296
0.554
D3
0.6096
5.35
0.267
0.622
D
0.0862
59.64
0.0935
0.945
init
D2
0.0785
57.58
0.1604
0.906
D3
0.0779
99.38
0.1509
0.929
17

---

Uncertainty-Base Extensible-Codecbook Federated Learning (UEFL)
Table 16: Our extensible codebook (Extend) outperforms the static larger codebook (Static) on all evaluation metrics.
Codebook
Data
Loode
mP
mE
mA
D1
3.10
17.54
1.66
0.142
Static
D2
2.91
17.41
1.76
0.135
D3
2.63
16.82
1.76
0.112
Di
1.28
19.31
0.7822
0.375
Extend
D2
0.976
27.42
0.6665
0.341
D3
0.978
22.00
0.7112
0.304
imposes stricter criteria, pushing the model to achieve higher precision, thereby improving performance metrics. However, it's important to recognize that beyond certain point, further reducing the threshold may not significantly enhance outcomes but will increase computational overhead. Thus, in such cases, there is trade-off between runtime and performance.
Performance of the 2nd distribution
Performance of the 3rd distribution
1.0
1.0
0.8
0.8
Metric (a.u.)
Metric (a.u.)
0.6
0.6
0.4
0.4
Accuracy Uncertainty Perplexity
Accuracy Uncertainty
0.2
0.2
Perplexity
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.0
0.1
0.2
0.3
0.4
0.5
0.6
Threshold
Threshold
Figure 9: Results on GTSRB dataset with 64 initialized codewords with segment 1. Overall, a smaller threshold performs better.
18


Reviews:
[{"review_id": "V4HxdMGTxc", "reviewer": "Reviewer_7WL8", "summary": "This paper introduces Uncertainty-Based Extensible-Codebook Federated Learning (UEFL), a novel framework addressing data heterogeneity in federated learning. The key idea is to dynamically extend a codebook of discrete latent vectors based on model uncertainty, which is evaluated via Monte Carlo Dropout. UEFL initializes a small shared codebook and iteratively adds client-specific codewords using K-means clustering on encoder features for underrepresented distributions. Experiments on rotated MNIST, CIFAR, GTSRB, and PACS datasets demonstrate improvements in accuracy and uncertainty reduction compared to FedAvg and DisTrans. The method also shows scalability to large client numbers and robustness in domain generalization tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the insightful comments and suggestions. We address your concerns as follows:\n\n> While the theorems ... connection to UEFL’s empirical success is not explicitly discussed\n\nWe provide a detailed analysis in response to the later theoretical questions. Please refer to that.\n\n> The paper mainly uses rotation transformation ... not cover more complex distribution drift situations\n\nOther forms of heterogeneity are also included:\n- Feature space shifts: Tables 2 and 3 (e.g., PACS domains)\n- Uneven class distributions: Table 12 (Label heterogeneity)\n\n> Lack of detailed discussion on the impact of different codebook sizes on model stability\n\nWe ran experiments across 5 random seeds. As the codebook size increases from 16 to 64, the standard deviation of accuracy drops from 0.0237 to 0.0094 (higher stability).\n\n> Only used 5 sub models ... not explore whether increasing the number of sub models would affect the stability of the assessment\n\nWe extended the ensemble size to 20 and observed the following on MNIST:\n| Method | Accuracy| Uncertainty|\n| ----- | ----- |----- |\n| FedAvg | 0.782 | 0.261|\n| UEFL | 0.924 | 0.237|\nUEFL also outperforms. However, consistent with prior findings (e.g., Lakshminarayanan et al., 2017), we observed that minimal improvement beyond 5–10 ensembles while computation costs increase linearly.\n\n> Appendix B ... not provide a detailed generalization error curve\n\nWe have plotted the curve: https://ibb.co/xymvPDT\n\n> Appendix F ... does not explore the impact of communication costs\n\nIn practice, UEFL uses 30 communication rounds in its first iteration, compared to 40 in standard FL. Each subsequent iteration requires only 5 rounds with a modest model size increase of 0.83%. For example, with two additional iterations, the overall communication cost is approximately $\\frac{30}{40}+\\frac{5}{40}\\times1.0083+\\frac{5}{40}\\times1.0167 = 1.0031\\times$ of standard FL.\n\n> More extensive experiments are still needed to verify their applicability\n\nWe’re extending this to medical data and other domains in future work.\n\n> Essential References Not Discussed\n\nThanks. Specifically for FedProx, the results are as follows:\n| Method| FMNIST|\n| -------- | ------- |\n| FedProx| 88.55   |\n| UEFL| 90.59    |\nwe will add more comparisons in the updated version.\n\n> The paper provides a theoretical analysis ... not consider the Non IID distribution in FL\n\nIn non-IID FL, then\n\n**With discretization:**\n\\begin{equation}\n   \\mathcal{O}( \\alpha \\sqrt{ \\frac{G \\ln L + \\ln(2K/\\delta)}{2n}} + \\frac{\\nu^{(q)}}{\\sqrt{n}}),\n\\end{equation}\n**Without discretization:**\n\\begin{equation}\n    \\mathcal{O}( \\alpha \\sqrt{ \\frac{m \\ln(4\\sqrt{nm}) + \\ln(2K/\\delta)}{2n} } + \\frac{\\overline{\\varsigma} R\\_\\mathcal{H} + \\nu}{\\sqrt{n}}),\n\\end{equation}\nHere, the KL divergence $\\nu^{(q)} < \\nu$.\nTherefore, discretization explicitly mitigates the effects of data heterogeneity.\n\nTo conserve space, more details are included in our response to Reviewer idfz.\n\n> Not provide a mathematical analysis of the effect of codebook size on convergence\n\nThe basic mathematical analysis follows VQ-VAE. Large codebooks hinder convergence due to poor utilization and noisy updates. To address this, UEFL progressively extends the codebook, initializing new codewords with K-means for better alignment.\nThis improves codeword utilization, reduces quantization error, and accelerates convergence (Fig. 3(b)).\n\n> Limited exploration of Label Heterogeneity\n\nWe further test different dirichlet distributions as follows,\n| Method| α=0.05| α=0.01| \n| ----- | ----- |----- |\n| FedAvg| 78.57   |35.42   |\n| UEFL|85.18    | 40.99   |\nUEFL performs better. More results will be included.\n\n> Only uses K-means for initialization ... not explore other possible initialization methods\n\nThanks. We agree that PCA or contrastive-based initialization is promising and will explore them in future work.\n\n> How does UEFL’s communication cost (e.g.,transmitting codeword vectors) scale with the number of clients and codebook size? For 100 clients (Table 4), does the server need to aggregate 100 unique codebooks?\n\nNo. Codebook growth is tied to iterations, not client count. In each iteration, 64 codewords are added and shared with selected high-uncertainty clients. The communication overhead after $i$ iterations is $Overhead(i) = \\frac{i\\times 0.125}{14.991}$.\nWith a typical $i=1-3$ and $i=5$ (Section 3.2), the worst-case overhead is under 4.17% (5 unique codebooks), regardless of client count (100).\n\n> Can Theorem 1 be extended to account for dynamic codebook growth?\n\nYes. Although Theorem 1 assumes static $L$, dynamic growth adds a minor $\\text{ln}L$ term. $L$ remains small (e.g., 64–256). However, dynamic codebook growth improves representation capacity, thereby reducing the KL divergence term $v^{(q)}$, tightening the federated generalization bound.\n\n> Experiment conducted under more extreme Non IID settings (α=0.01)?\n\nThe results for Dirichlet distribution α=0.01 are discussed above.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Theorems 1 and 2 in Appendix A argue that discretization reduces generalization gaps by lowering noise sensitivity and dimensionality. While the theorems are logically structured, their connection to UEFL’s empirical success is not explicitly discussed. For instance, how the codebook’s extensibility interacts with the theoretical guarantees remains unclear. The proof relies on Hoeffding's Inequality, but does not explicitly discuss which distribution data the inequality applies to, especially in the Non IID scenario of federated learning, which may affect its applicability.", "experimental_designs_or_analyses": "Limited simulation of data heterogeneity: The paper mainly uses rotation transformation to introduce heterogeneity, but in real federated learning environments, data heterogeneity usually includes uneven class distribution, feature space shift, etc., and does not cover more complex distribution drift situations.\n\nLack of generalization analysis for different codebook sizes: Although the paper provides some experiments with different codebook sizes \nsuch as K-means initialization vs. random initialization), there is a lack of detailed discussion on the impact of different codebook sizes on model stability.\n\nInsufficient ablation experiments for uncertainty assessment methods: Although the paper compared Deep Ensemble, the experiment only used 5 sub models and did not explore whether increasing the number of sub models would affect the stability of the assessment.", "relation_to_broader_scientific_literature": "The core contributions of this paper mainly involve two fields: Federated Learning (FL) and Uncertainty Modeling, and are related to existing research as follows:\n\nData heterogeneity is a key challenge in FL, and various methods have been proposed in previous studies to alleviate this issue: Personalized FL based methods (such as FedPer, FedRep): Processing heterogeneous data through hierarchical separation or personalized local models [Li et al., 2021]. FL methods based on distribution transformation (such as DisTrans): Processing data heterogeneity through distribution transformation during training and testing [Yuan et al., 2022]. FL methods based on knowledge distillation (such as FCCL): using knowledge distillation and unlabeled public data to enhance generalization ability [Huang et al., 2022]. The core innovation of UEFL in this article lies in the extensible codebook, which is similar to the idea of enhancing FL robustness through discretization (such as VQ FedAvg [Liu et al., 2021]). However, this article additionally introduces a dynamic codebook extension based on uncertainty to adapt to data with larger distribution biases.\n\nPrevious studies have used uncertainty quantification to improve the robustness of FL models: Method based on Monte Carlo Dropout (Gal&Ghahramani, 2016): Estimating uncertainty by enabling Dropout during the inference phase for data selection and model weighting. Method based on Deep Ensemble (Lakshminarayanan et al., 2017): Train multiple models and estimate uncertainty through analysis of variance. This article uses Monte Carlo Dropout as an uncertainty assessment method and combines it with K-means to dynamically extend the codebook, further expanding the application of uncertainty modeling in the field of FL.\n\nThe paper mentions VQ-VAE [Van Den Oord et al., 2017] as inspiration for using a discrete codebook for feature mapping. In the field of FL, Liu et al.,  The use of discretization to enhance model generalization ability has been proposed in 2021, but the encoding of this method is fixed. This article proposes an extensible codebook mechanism and dynamically adjusts the codebook size by combining uncertainty, which is one of the main innovations of this article.\n\nIn summary, the innovation of this paper mainly lies in the combination of uncertainty quantification and scalable codebook mechanism to handle FL data heterogeneity, and empirical research on multiple datasets. These contributions complement existing literature, but more extensive experiments are still needed to verify their applicability (such as other data types, communication overhead, etc.).", "essential_references_not_discussed": "VQ-FL (Chen et al., 2023, ICML): Uses vector quantization for client-specific representation learning but fixes the codebook size. UEFL’s uncertainty-driven extension is novel, but a comparison is necessary.\n\nFedPM (Dinh et al., 2022, NeurIPS): Personalizes models via latent mask vectors. While distinct from codebooks, its focus on client-specific latent spaces is conceptually related.\n\nFedProx (Li et al., 2020, MLSys): Addresses heterogeneity via proximal regularization. Although cited in the related work, its comparison with UEFL in terms of uncertainty reduction is missing.", "other_strengths_and_weaknesses": "Strengths:\n\nA dynamic codebook extension strategy based on uncertainty assessment has been proposed, which is more adaptable compared to existing discretization FL methods such as VQ FedAvg.\n\nUsing K-means for codebook initialization reduces the instability caused by random initialization and improves model training efficiency.\n\nMultiple datasets (MNIST, FMNIST, CIFAR10, CIFAR100, GTSRB) were used for experiments, and rotation transformation was employed to simulate data heterogeneity.\n\nWeaknesses:\n\nThe paper provides a theoretical analysis of the discretization of generalization error in Appendix A, but this analysis relies on the i.i.d. assumption and does not consider the Non IID distribution in FL.\n\nThe paper does not provide a mathematical analysis of the effect of codebook size on convergence, and only verifies its impact through experiments.\n\nThe paper mainly simulates Feature Heterogeneity, but there is limited exploration of Label Heterogeneity, with only limited experiments conducted in Appendix H.\n\nThe paper only uses K-means for initialization of the codebook and does not explore other possible initialization methods (such as PCA dimensionality reduction and contrastive learning feature initialization).", "other_comments_or_suggestions": "See the weakness", "questions_for_authors": "Communication Overhead: How does UEFL’s communication cost (e.g., transmitting codeword vectors) scale with the number of clients and codebook size? For 100 clients (Table 4), does the server need to aggregate 100 unique codebooks?\n\nTheoretical Grounding: Can Theorem 1 be extended to account for dynamic codebook growth? For example, does adding codewords tighten the generalization bound in Eq. 8?\n\nThe mathematical analysis of the paper assumes that the data is i.i.d., but FL is usually a Non IID scenario. Is there an experiment conducted under more extreme Non IID settings (such as Dirichlet distribution α=0.01)? Is the scalability of UEFL still effective for situations where data distribution is severely uneven across different clients?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wlGJs0icWz", "reviewer": "Reviewer_idfz", "summary": "The paper introduces Uncertainty-Based Extensible-Codebook Federated Learning to address data heterogeneity in federated learning (FL) by dynamically expanding a discrete codebook based on model uncertainty. UEFL improves generalization by mapping latent features to trainable codewords and selectively extending the codebook for clients exhibiting high uncertainty, reducing performance degradation caused by non-IID data distributions. The approach integrates Monte Carlo Dropout for uncertainty evaluation and K-means clustering for efficient codeword initialization, ensuring minimal computational overhead.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the insightful comments and suggestions. In response, we address these issues below:\n\n> Scalability to thousands of clients is not tested. Estimate overhead w.r.t. the number of clients\n\nWhile we did not test UEFL with thousands of clients, we evaluated it with 50 & 100 clients, showing consistent improvements over baselines (Tables 3 & 4). UEFL’s overhead scales with iteration count, not the number of clients. In each iteration, 64 codewords are added and shared with all selected high-uncertainty clients; the overhead after $i$ iterations is:\n\n\\begin{equation}\n    Overhead(i) = \\frac{i\\times 0.125}{14.991}\\times 100\\\\%\n\\end{equation}\n\nUEFL typically needs 1-3 iterations and we set the maximum to 5 (Section 3.2), so the introduced overhead is at most 4.17%.\n\n> In some FL settings, stricter constraints (e.g., differential privacy, homomorphic encryption) exist\n\nThanks for the comment. Our current experimental setup assumes standard FL settings, but UEFL is compatible with these privacy-preserving methods. Specifically, our codeword-based discretization process can integrate differential privacy by adding calibrated noise to codeword embeddings. We will explore this in future work.\n\n> The theoretical section ... rely on IID assumptions, while real FL data is often non-IID\n\nIn non-IID FL with $K$ clients, client $k$ has $n_k$ samples drawn from its distribution $P_k$. Total samples $n = \\sum_{k=1}^Kn_k$. The global distribution is $\\overline{P} = \\sum_{k=1}^K \\frac{n_k}{n}P_k$, then\n\n**With discretization:**\n\n\\begin{equation}\n    \\left|\\sum_{k=1}^K \\frac{n_k}{n} \\mathbb{E}_{\\boldsymbol{h} \\sim P_k}[\\phi_k^S(q(\\boldsymbol{h}, L, G))] - \\frac{1}{n} \\sum\\_{k=1}^K \\sum\\_{i=1}^{n_k} \\phi_k^S(q(\\boldsymbol{h}_i^{(k)}, L, G))\\right| = \\mathcal{O}( \\alpha \\sqrt{ \\frac{G \\ln L + \\ln(2K/\\delta)}{2n}} + \\frac{\\nu^{(q)}}{\\sqrt{n}}),\n\\end{equation}\n\n**Without discretization:**\n\n\\begin{equation}\n    \\left|\\sum_{k=1}^K\\frac{n_k}{n} \\mathbb{E}_{\\boldsymbol{h} \\sim P_k}[\\phi_k^S(\\boldsymbol{h})] - \\frac{1}{n} \\sum\\_{k=1}^K \\sum\\_{i=1}^{n_k} \\phi_k^S(\\boldsymbol{h}_i^{(k)}) \\right| = \\mathcal{O}( \\alpha \\sqrt{ \\frac{m \\ln(4\\sqrt{nm}) + \\ln(2K/\\delta)}{2n} } + \\frac{\\overline{\\varsigma} R\\_\\mathcal{H} + \\nu}{\\sqrt{n}}),\n\\end{equation}\n\nHere, $\\nu^{(q)} = \\frac{1}{K}\\sum_{k=1}^K\\text{Div}(P_k^{(q)}, \\overline{P}^{(q)})$ and $\\nu = \\frac{1}{K}\\sum_{k=1}^K\\text{Div}(P_k, \\overline{P})$ denote the KL divergence between client and global distribution with and without discretization. $\\nu^{(q)} < \\nu$.\nTherefore, discretization not only improves robustness to noise and reduces dimensionality, but also it effectively mitigates the effects of data heterogeneity typical in non-IID FL.\n\n> Figure 9 ... no principled way to set threshold\n\nIn this work, the threshold is manually tuned per dataset. We agree that a dynamic adjustment mechanism (e.g., based on convergence) is promising future work.\n\n> Not analyze long-term codebook growth ... analysis of the size of the codebook on the uncertainty\n\nComputation is discussed above.\n\nCodeword utilization, measured by perplexity ($\\exp(-\\sum_{\\text{class}} p \\log p)$), initially increases as the codebook expands, leading to richer representations and higher mutual information $I(Z; Y)$, which reduces uncertainty ($H(Y|Z) = H(Y) - I(Z; Y)$). However, when the codebook becomes very large, most codewords are rarely used (collapse) [1], resulting in reduced perplexity and diminishing returns in uncertainty reduction. More details will be included.\n\n> The plots are blurry. Use vector illustration\n\nThanks. We have replaced all figures with vector graphics to ensure clarity. Please review a few updated samples (.pdf figures): https://blush-melessa-85.tiiny.site\n\n> The running title needs an update.\n\nWe have updated it to: \"UEFL: Uncertainty-Based Extensible Codebook Federated Learning\".\n\n> The experimental setting is unclear ... Use both GPUs or only one?\n\nAll experiments utilized only one GPU, added to the revision.\n\n> How does performance degrade if the codebook is not expanded enough?\n\nThen, the model lacks representational capacity, resulting in reduced accuracy and higher uncertainty (Figure 6). However, UEFL typically converged after 1-3 iterations.\n\n> Incorporate newer techniques like FedPer, FedRod, FedBABU, or FedDyn\n\nHere is the comparison with FedRod & FedDyn:\n| Method| FMNIST|\n| -------- | ------- |\n| FedDyn| 89.65   |\n| FedRod| 90.28   |\n| UEFL| 90.59    |\n\nUEFL outperforms them. More results will be included in the updated version.\n\n> This work utilizes a pre-trained VGG model ... justify the necessity of the proposed approach, given that the model is already well-suited for these datasets?\n\nWhile the encoder is strong, the classifier is still vulnerable to domain shifts. UEFL improves performance via discretization, as shown in Table 1.\n\n[1] Huh et al., Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks, ICML 2023.", "claims_and_evidence": "The computational overhead of UEFL is minimal and does not impact scalability. While UEFL’s memory overhead is shown to be small (~3.34% increase), scalability to thousands of clients is not tested. It is also useful to provide the estimation of the overhead w.r.t. the number of clients.", "methods_and_evaluation_criteria": "FL often deals with privacy-preserving scenarios where clients cannot share model updates freely. The paper assumes all clients can exchange information, but in some FL settings, stricter constraints (e.g., differential privacy, homomorphic encryption) exist.", "theoretical_claims": "The theoretical section (Appendix A) provides proofs justifying discretization in FL. However, they rely on IID assumptions, while real FL data is often non-IID.", "experimental_designs_or_analyses": "Figure 9 shows that lower thresholds improve performance, but there is no principled way to set the threshold.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The paper does not analyze long-term codebook growth, which could become a computational bottleneck over extended training.\n\nThis work lacks theoretical support, e.g., analysis of the size of the codebook on the uncertainty.", "other_comments_or_suggestions": "1. The plots are blurry when I zoom in. It is better to use vector illustration.\n2. The running title needs an update.\n3. The experimental setting is unclear. It says the \"experiments are performed on a machine with 2 GPUs\". Did you use both GPUs or only one GPU?", "questions_for_authors": "1. How does performance degrade if the codebook is not expanded enough?\n\n2. The baseline comparison primarily includes FedAvg (2017) and DisTrans (2022), which are outdated given recent advancements in federated personalization and adaptive aggregation methods; incorporating newer techniques like FedPer, FedRod, FedBABU, or FedDyn would provide a more rigorous evaluation of UEFL’s effectiveness.\n\n3. This work utilizes a pre-trained VGG model for tasks like CIFAR and GTSRB, which already have strong feature extractors. How does this justify the necessity of the proposed approach, given that the model is already well-suited for these datasets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H8EHdCv1uV", "reviewer": "Reviewer_t3oA", "summary": "The paper addresses the challenge of data heterogeneity in federated learning (FL) by proposing Uncertainty-Based Extensible-Codebook Federated Learning (UEFL). The method dynamically extends a codebook of latent vectors using uncertainty estimates (via Monte Carlo Dropout) to adapt to diverse data distributions across silos. Key innovations include K-means initialization for new codewords, segmentation of feature vectors, and iterative codebook expansion. Experiments on rotated datasets (MNIST, CIFAR, etc.) demonstrate improvements in accuracy (3%-22.1%) and uncertainty reduction (38.83%-96.24%) over FedAvg and DisTrans. The approach also scales well to large client numbers (50–100) and handles domain generalization tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for the insightful and constructive feedback, as well as the recognition of our contributions and experiments. Below, we address the specific questions raised:\n\n> What is the extra information, except for the model parameters, that the clients send to the central server? Is there any risk of information leakage during the codebook sharing?\n\nIn our UEFL framework, the codebook is a trainable component of the model architecture (like classifier weights or encoder layers) and is thus included in the standard model parameters exchanged during federated averaging. Clients only share updated model parameters (including discrete codebook vectors) with the server, no raw data or additional metadata.\n\nCodebooks map latent features to discrete codewords, which represent aggregated and abstracted latent features rather than raw data, less prone to information leakage. Discrete representations inherently limit the granularity of shared information, aligning with privacy-preserving mechanisms in federated learning [1, 2]. In addition, the segmentation of codewords (Section 3.2) further abstracts feature information, enhancing the robustness against potential information leakage compared to raw data or explicit feature representations. \n\nWe will clearly clarify these points in the revised manuscript to explicitly address potential privacy concerns. We deeply appreciate the thoughtful review and suggestions.\n\n[1] Kairouz, Peter, Ziyu Liu, and Thomas Steinke. \"The distributed discrete gaussian mechanism for federated learning with secure aggregation.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Agarwal, Naman, Peter Kairouz, and Ziyu Liu. \"The skellam mechanism for differentially private federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 5052-5064.", "claims_and_evidence": "The claims in the paper are supported by clear evidence:\n1. The proposed method shows consistent improvement in MNIST, PACS, CIFAR10, and CIFAR 100. The authors have designed sufficient experiments to prove the effectiveness of the proposed UEFL. \n2. The authors have conducted careful ablation studies on the hyperparameters (e.g. Uncertainty Threshold)\n3. The paper also gives a theoretical explanation in Appendix A.", "methods_and_evaluation_criteria": "Yes, the proposed method is simple and straightforward in solving the problem of limited codebooks when heterogeneous data occurs. It is expected to get an improvement. The experiments also prove the method's effectiveness.", "theoretical_claims": "There is no proof in the paper.", "experimental_designs_or_analyses": "The paper have conducted a series of experiments on MNIST, PACS, CIFAR10, and CIFAR 100. Careful ablation studies have been done.", "relation_to_broader_scientific_literature": "The paper is related to privacy-preserving training.", "essential_references_not_discussed": "The authors discussed the related works well.", "other_strengths_and_weaknesses": "Strength:\n1. The algorithm introduces light overhead, which is suitable for edge deployment.\n2. The paper is well written and easy to read.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. What is the extra information, except for the model parameters, that the clients send to the central server? Is there any risk of information leakage during the codebook sharing?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f36Lur09bv", "reviewer": "Reviewer_SFNQ", "summary": "This paper introduces Uncertainty-Based Extensible-Codebook Federated Learning (UEFL), a novel framework designed to address data heterogeneity in federated learning (FL). The key innovation lies in dynamically mapping latent features to trainable discrete vectors (codewords) and extending the codebook for silos with high uncertainty, identified via Monte Carlo Dropout. UEFL demonstrates significant improvements in accuracy and uncertainty reduction across various datasets, including MNIST, CIFAR10, and CIFAR100. The extensible codebook approach, initialized using K-means, ensures efficient adaptation to unseen data distributions while maintaining low computational overhead.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the recognition of the novelty of our method, and the robustness of our experimental design. Additionally, we value the insightful critique regarding the limitations of our work. In response, we address these issues below:\n\n> However, the introduction of an extensible codebook raises potential concerns regarding privacy risks that warrant further discussion. Since the codebook is shared and updated across clients, there may be a risk of leakage of sensitive information embedded in the codewords, especially if adversaries attempt to reverse-engineer the mapping between latent features and codewords.\n\nWe thank you for highlighting the important concern regarding potential privacy risks associated with the extensible codebook. To ensure robust privacy preservation, our method follows standard federated learning protocols, raw data never leaves the local client, and only aggregated model updates are communicated. The codewords represent abstract latent features derived from model encoders, rather than explicit raw data or identifiable content. These latent vectors are highly compressed and abstract, significantly reducing the feasibility of reverse-engineering meaningful private information. \n\nFurthermore, the segmentation of codewords further abstracts feature information. Newly added codewords for highly heterogeneous silos are client-specific and accessible exclusively by those clients, minimizing potential information leakage risks. We will explicitly discuss these privacy aspects in our updated manuscript to highlight our framework's robust measures against potential privacy risks.\n\nNonetheless, we acknowledge the insightful suggestion and completely agree that the introduction of an extensible codebook could raise potential privacy risks. However, this is not the focus of this work, as our primary aim was to address model accuracy and uncertainty reduction in heterogeneous data silos within federated learning settings. Possible solutions to address privacy concerns include differential privacy [1] or secure aggregation [2].\nWe will leave the detailed exploration and integration of such privacy-preserving techniques for further study.\n\n\n> While this paper emphasizes the alignment of codewords with latent features via K-means initialization to improve model performance, it remains unclear how this process is safeguarded against potential attacks, such as model inversion or membership inference attacks.\n\nWe appreciate this insightful comment regarding the robustness of our K-means initialization method against potential attacks. Our K-means algorithm operates solely on client-specific, highly abstracted latent embeddings rather than raw data, inherently limiting the risk of reconstructing original inputs. Additionally, our approach is inherently compatible with advanced privacy-preserving techniques, such as differential privacy [1], where calibrated noise can be added to the discrete codeword embeddings to further enhance security.\n\nAlthough our primary goal was to address model accuracy and uncertainty reduction, we fully acknowledge the concerns and will explicitly mention this security consideration in the revised manuscript, noting this as an area for further detailed exploration.\n\n> While the provided theorems establish the foundational advantages, including enhanced robustness to noise and reduced dimensionality, a more detailed step-by-step proof process would further strengthen the theoretical claims and improve the clarity of the mathematical reasoning.\n\nWhile Appendix A provides theoretical foundations supporting the benefits of discretization, we agree that providing a detailed, step-by-step derivation of our theoretical results would strengthen the clarity of our claims.\nIn the revised supplementary materials, we will include full derivations based on Hoeffding Inequality for Theorem 1 and 2.\n\n> The resolution of the figures in the paper is quite low, and some of the text within the images is difficult to read due to its small size.\n\nThank you for highlighting this issue. We have improved the readability and resolution of our figures by regenerating them using vector graphics in PDF format, with increased font sizes for better clarity. Please review a few updated figure samples (.pdf figures) at the following anonymous link: https://blush-melessa-85.tiiny.site\n\nThanks once again for the constructive comments and valuable suggestions, which have significantly enhanced the quality and clarity of our manuscript.\n\n[1] Agarwal, Naman, Peter Kairouz, and Ziyu Liu. \"The skellam mechanism for differentially private federated learning.\" Advances in Neural Information Processing Systems 34 (2021): 5052-5064.\n\n[2] Kairouz, Peter, Ziyu Liu, and Thomas Steinke. \"The distributed discrete gaussian mechanism for federated learning with secure aggregation.\" International Conference on Machine Learning. PMLR, 2021.", "claims_and_evidence": "The claims of improved accuracy and uncertainty reduction are well-supported by empirical evidence from experiments on multiple datasets.", "methods_and_evaluation_criteria": "The proposed method is well-suited to the problem of feature heterogeneity in FL, and the use of an extensible codebook with uncertainty evaluation is innovative. However, the introduction of an extensible codebook raises potential concerns regarding privacy risks that warrant further discussion. Since the codebook is shared and updated across clients, there may be a risk of leakage of sensitive information embedded in the codewords, especially if adversaries attempt to reverse-engineer the mapping between latent features and codewords. \n\nWhile this paper emphasizes the alignment of codewords with latent features via K-means initialization to improve model performance, it remains unclear how this process is safeguarded against potential attacks, such as model inversion or membership inference attacks.", "theoretical_claims": "This paper provides theoretical support for the benefits of discretization in reducing noise sensitivity and dimensionality, as outlined in Appendix A. While the provided theorems establish the foundational advantages, including enhanced robustness to noise and reduced dimensionality, a more detailed step-by-step proof process would further strengthen the theoretical claims and improve the clarity of the mathematical reasoning.", "experimental_designs_or_analyses": "The experimental design is robust, with ablation studies and comparisons against strong baselines like FedAvg and DisTrans. The use of domain generalization tasks and large-scale setups (e.g., 50 and 100 clients) further strengthens the evaluation.", "relation_to_broader_scientific_literature": "UEFL builds on prior work in FL, such as FedAvg and DisTrans, while introducing a novel mechanism to handle data heterogeneity through extensible codebooks and uncertainty-based adaptation. It also aligns with recent trends in uncertainty modeling (e.g., Monte Carlo Dropout, Deep Ensembles) and discrete representations for robustness.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "The resolution of the figures in the paper is quite low, and some of the text within the images is difficult to read due to its small size.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper has addressed the data distribution shift in discrete federated learning. The core idea is to dynamically extend a codebook of latent vectors using uncertainty estimates, i.e., Monte Carlo Dropout, to adapt to diverse data distributions across silos. The problem is well defined, and the method has been well formulated. Some good results have been made.

Some suggestions, i.e.,

- Include the experiments added in rebuttal in the final version.
- Revise some figures with low resolution.
Decision: Accept (poster)