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

## Abstract

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 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. Through experiments conducted on six datasets, our method has demonstrated its superiority, achieving significant improvements in accuracy (by 3%–22.1%) and uncertainty reduction (by 38.83%–96.24%), thereby outperforming contemporary state-of-the-art methods.

## 1 Introduction

Federated Learning (FL), well known for its capacity to harness data from diverse devices and locations—termed data silos—while ensuring privacy, has become increasingly crucial in the digital era, particularly with the explosion of data from mobile sources.
Despite its pivotal role in distributed computing, FL confronts a formidable challenge: the heterogeneity of data across different silos. Such diversity often results in a significant performance gap when integrating updates from local models into the global model. In Fig. 1, we 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 on varied data sources.

Recent studies (Ghosh et al., ; Agarwal et al., ; Liu et al., ; Kairouz et al., ; Zhang et al., ; Yuan et al., ) have made significant advancements in addressing data heterogeneity within Federated Learning (FL), with one notable approach being the use of discrete representations to enhance model robustness against minor data shifts. Nonetheless, this strategy struggles to generalize models to data silos exhibiting significant distributional differences. Furthermore, these methods face difficulties in adapting to unseen data distributions, as they typically require the entire model to be re-trained. Such constraints limit their flexibility in adapting to the dynamically changing data landscapes, posing challenges for their applicability in real-world scenarios.

Figure: Figure 1: In the case of heterogeneous data silos, the global model of regular Federated Learning (FL) performs poorly compared to local models before integration. By discretizing different domains into distinct latent spaces, our UEFL improves both accuracy and uncertainty. The reported values in the figure represent the average accuracy and uncertainty across the various data silos.
Refer to caption: extracted/5973906/intro_iclr.png

Moreover, we identify another critical issue impacting the model’s performance across diverse data silos: increased uncertainty, as shown in Fig. 1. The global model’s accuracy not only deteriorates, but its uncertainty also trends upwards, signaling increased prediction instability. To address these challenges, we introduce Uncertainty-Based Extensible-codebook Federated Learning (UEFL), a novel methodology that explicitly distinguishes between data distributions to improve both accuracy and uncertainty.

Specifically, our design features an advanced codebook comprising a predetermined number of latent vectors (*i.e*. codewords), and employs a discretizer to assign encoded image features to their closest codewords. These codewords, acting as latent representations, are passed to 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 initialize a small, shared codebook for all clients. Additional specific codewords are then introduced for individual client use, ensuring explicit differentiation between them. Since the initial codebook is small and requires only a few extensions, the final size remains compact, minimizing the associated computational overhead.
Given the privacy constraints in federated learning (FL), which restrict direct data access, we incorporate an uncertainty evaluator using Monte Carlo Dropout. This evaluator identifies data from diverse distributions, marked by high uncertainty. During training, our UEFL method systematically distinguishes 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 leveraged to initialize new codewords using K-means, aligning them more closely with the data distribution and facilitating faster adaptation to 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. Furthermore, since uncertainty constantly decreases as training progresses, the iterative process is guaranteed to conclude after a few iterations.

To summarize, our contributions are as follows:

- •
We identify a significant increase in model uncertainty across silos with diverse data distributions within the federated learning (FL) context, highlighting the challenge of data heterogeneity.
- •
To address this heterogeneity, we introduce an extensible codebook approach that distinguishes between data distributions by stepwise mapping them to distinct, trainable latent vectors (*i.e*. codewords). This methodology allows for efficient initialization of newly added codewords using a K-means algorithm, closely aligning with the training data feature distributions and enabling rapid convergence during codebook training.
- •
We propose a novel data-driven FL framework, named Uncertainty-Based Extensible-codebook Federated Learning (UEFL), which merges the extensible codebook with an uncertainty evaluator. This framework iteratively identifies data from diverse distributions by assessing uncertainty without requiring direct data access. It then processes this data by initializing new codewords to complement the existing codebook, ensuring that each iteration focuses on training the expandable codebook, which rapidly converges, thus allowing UEFL to adapt seamlessly to new data distributions.
- •
Our empirical evaluation across six datasets demonstrates that our approach significantly reduces uncertainty by 38.83%-96.24% and enhances model accuracy by 3%-22.1%, evidencing the effectiveness of UEFL in managing data heterogeneity in FL.

## 2 Related Work

### 2.1 Federated Learning

Federated learning (Konečnỳ et al., ; Geyer et al., ; Chen et al., ; Hard et al., ; Yang et al., ; Ghosh et al., ) represents a cutting-edge distributed learning paradigm, specifically designed to exploit data and computational resources across edge devices. The Federated Averaging (FedAvg) algorithm (McMahan et al., ), introduced to address the challenges of unbalanced and non-IID data, optimizes the trade-off between computation and communication costs by reducing the necessary communication rounds for training deep networks. Federated Learning (FL) faces numerous statistical challenges, with data heterogeneity being one of the most critical. In real-world applications, data collected across different clients often varies significantly in terms of distribution, feature space, and sample sizes.

Several methodologies (Zhao et al., ; Li et al., ; ; Kalra et al., ) have been developed to address this pivotal issue. PMFL (Zhang et al., ) approaches the heterogeneity challenge by drawing inspiration from meta-learning and continual learning, opting to integrate losses from local models over the aggregation of gradients or parameters. DisTrans (Yuan et al., ) enhances FL performance through train and test-time distributional transformations, coupled with a novel double-input-channel model architecture. Meanwhile, FCCL (Huang et al., ) employs knowledge distillation during local updates to facilitate 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 heterogeneity by Liu et al. (), provide further inspiration and valuable perspectives for our research endeavors.

### 2.2 Uncertainty

Recently, the study of uncertainty modeling has gained significant prominence across various research fields, notably within the machine learning community (Chen et al., ; Blundell et al., ; Kendall & Gal, ; Louizos & Welling, ; Lahlou et al., ; Nado et al., ; Gawlikowski et al., ). 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, ), which introduces variability in model outputs through the use of dropout layers, and Deep Ensembles (Lakshminarayanan et al., ), which leverages multiple models with randomly initialized weights trained on identical datasets to evaluate uncertainty, exemplify the advancements in this area. Furthermore, the application of uncertainty modeling has extended beyond traditional domains, impacting fields such as healthcare (Dusenberry et al., ) and continual learning (Ahn et al., ).

## 3 Methodology

### 3.1 Overall Architecture

Figure: Figure 2: UEFL flowchart. In the first iteration, all latent are mapped to initialized shared codewords by the discretizer $\theta_{D}$. 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.
Refer to caption: extracted/5973906/chart.png

Fig. 2 illustrates the workflow of our UEFL.
Consider multiple data distributions $\mathcal{D}_{1},\mathcal{D}_{2},...,\mathcal{D}_{M}$, with data samples $x\in\mathbb{R}^{H\times W\times D}$, where $H$, $W$, and $D$ denote the input image’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 $\theta_{E}$ into feature representations $z\in\mathbb{R}^{h\times w\times d}$, with $h$, $w$, and $d$ representing the features’ shape. Subsequently, these features are reshaped into vectors $z\in\mathbb{R}^{l\times d}$, where $l$ is the number of tokens, and divided into $s$ segments $z_{i}\in\mathbb{R}^{l\times\frac{d}{s}},\forall i$, with $s$ indicating the segment count. Each segment is mapped to the closest codeword in the codebook via a discretizer $\theta_{D}$, then reassembled into complete vectors for classification. The classifier $\theta_{C}$ then deduces the class for the input data, completing the forward processing sequence as follows,

$$ $z=f_{\theta_{E}}(x),\;\;\;c=f_{\theta_{D}}(z),\;\;\;p=f_{\theta_{C}}(c)$ (1) $$

where $x$, $z$, $c$, and $p$ denote input data, latent features, discrete coded vectors, and the model prediction, separately.

After loss calculation, models undergo local updates through backpropagation. In a manner akin to FedAvg (McMahan et al., ), these updated models are then relayed back to the server for a global update.

$$ $\theta_{k}\leftarrow\theta-\eta g_{k},\;\forall k$ (2) $$

$$ $\theta\leftarrow\sum_{k=1}^{K}\frac{n_{k}}{n}\theta_{k},$ (3) $$

where $\theta$ denots the global model parameters, $\theta_{k}$ is the $k$th local model parameters, $g_{k}$ is the $k$th model gradients, $n_{k}$ is the number of samples for data silo $k$, and $n$ is the total number of samples for all $K$ silos.

At the end of each iteration, assessing uncertainty through Monte Carlo Dropout is essential, given the privacy constraints of Federated Learning (FL), which limit direct access to client data. By evaluating uncertainty against a pre-established threshold, we identify data from heterogeneous distributions. When such data are detected, we augment the codebook with new codewords, initializing them using centroids of the generated image features, as described in Algorithm 1. This process leverages the fully adapted encoder 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 better after integration, effectively handling data heterogeneity.

Figure: Algorithm 1 Uncertainty-Based Extensible-Codebook Federated Learning (UEFL)

### 3.2 Extensible Codebook

To effectively manage heterogeneous data, we design an extensible codebook, beginning with a minimal set of codewords and progressively enlarging this set through a superior initialization strategy that benefits from our UEFL framework. This strategy facilitates stepwise mapping of diverse data distributions to distinct codewords. Starting with a larger codebook can introduce uncertainty in codeword selection due to the concurrent training of multiple codewords.

Similar to VQ-VAE (Van Den Oord et al., ), we employ latent vectors as codewords, initializing a compact shared codebook with $n$ codewords $c\in\mathbb{R}^{n\times\frac{d}{s}}$, where $n$ represents the size of the initial codebook. The codewords are initialized using a Gaussian distribution and shared across all data silos. After each iteration’s uncertainty assessment, we determine which silos require additional codewords to improve prediction accuracy, and we extend the codebook accordingly for these silos by adding $n$ more codewords..
The newly added codewords are initialized using K-means rather than a Gaussian distribution, leveraging the encoder’s improved latent features from the prior iteration to better align with the underlying data distribution.
To optimize codebook usage, data silos that demonstrated lower performance in the previous iteration are allowed to select codewords from both the newly added codewords and the original shared codebook. Typically, the codebook only requires 1-3 extensions until all clients reach low uncertainty levels.

For a given iteration, if the codebook size for the $k$th client is $n_{k}$, the feature vector $z$ is associated with a codeword $c_{i}$ by the discretizer, which computes the distance between $z$ and all available codewords, selecting the nearest one as follows,

$$ $i=\operatorname*{arg\,min}_{j\in{1,2,...,n_{k}}}||z-c_{j}||_{2}$ (4) $$

K-means Initialization.
After the first iteration, the adapted encoder produces image features that more accurately reflect the distribution of the training data. Instead of relying on random initialization methods like Gaussian distribution, we initialize new codewords using the centroids of these features, obtained through K-means clustering. This approach expedites codebook training by providing a more informed starting point for the new codewords, allowing them to better align with the underlying data structure. As a result, this initialization strategy facilitates faster convergence and improves the model’s ability to adapt to varying data distributions across silos.
This strategy hugely reduces the number of training rounds required for model convergence in subsequent iterations (Details in Appendix B).

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 methodology, we dissect features into smaller segments to pair them with multiple codewords, thus covering the entirety of a feature vector.
This segmentation exponentially increases the codeword pool, ensuring a robust representation capacity without necessitating a large-scale increase and permitting efficient K-means-based initialization. This design minimizes runtime overhead associated with larger codebooks.

### 3.3 Loss Function

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 as follows:

$$ $\mathcal{L}_{task}=-\sum_{class}y\log p$ (5) $$

where $p$ is the model output, and $y$ is the ground truth label.

For codebook optimization, akin to the strategy employed in VQ-VAE (Van Den Oord et al., ), we apply a stop gradient operation for the codeword update as follows:

$$ $\mathcal{L}_{code}=||sg(c)-z||^{2}_{2}+\beta||c-sg(z)||^{2}_{2}$ (6) $$

where $z$ is the image latent features, $c$ is discrete codewords, $\beta$ is a hyper-parameter to adjust the weights of two losses and $sg(\cdot)$ denotes the stop gradient function.

So, the total loss $\mathcal{L}_{UEFL}$ is the summation of $\mathcal{L}_{task}$ and $\mathcal{L}_{code}$.

### 3.4 Uncertainty Evaluation

As outlined in Section 3.1, 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, ) 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 a variety of outcomes for uncertainty analysis.
This variability is quantified using predictive entropy, as described in Eq. 7, which serves to measure the prediction dispersion across different evaluations effectively.

$$ $e=-\sum_{class}p\log p$ (7) $$

A low predictive entropy value signifies model confidence, whereas a high value indicates increased uncertainty.
For high entropy, introducing new codewords and conducting additional training rounds are essential steps. Given the variability of uncertainty across datasets, establishing a fixed threshold is impractical. Instead, by analyzing all uncertainty values, we can benchmark against either the minimum or mean values to pinpoint target silos. Our experiments showed superior results when using the minimum value as a reference, thus guiding us to adopt the following threshold criterion:

$$ $e_{k}\leq(1+\gamma)\min_{\forall j\in 1,2,...,K}(e_{j}),\forall k$ (8) $$

where $\gamma$ is a hyperparameter to be set.

Furthermore, since uncertainty consistently decreases as training progresses—unlike accuracy, which may plateau after a certain point—the iterative process is guaranteed to conclude within a few iterations by evaluating uncertainty.

## 4 Experimental Results

Experimental Setup.
As discussed in Kairouz et al. (); Zhou et al. (), there are two predominant forms of data heterogeneity in federated learning: feature heterogeneity and label heterogeneity.
Our UEFL focuses on tacking feature heterogeneity, and we mainly discuss feature heterogeneity in this section. The discussion for label heterogeneity with dirichlet distribution and the comparison with VHL (Tang et al., ), FedBR (Guo et al., ) are in the Appendix A.

Similar to Rotated MNIST (Ghifary et al., ), which creates 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 different datasets: MNIST, FMNIST, CIFAR10, GTSRB, and CIFAR100, to validate our framework’s robustness.
In our experiments, we create three domains by counter-clockwise rotating the datasets by 0° ($\mathcal{D}_{1}$), -50° ($\mathcal{D}_{2}$), and 120° ($\mathcal{D}_{3}$). We sampled three data silos from each domain (*i.e*. totally 9 silos), and data silos for CIFAR100 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 out UEFL for domain generalization (DG) task on Rotated MNIST (Ghifary et al., ) and PACS (Li et al., ) datasets, which contains four distinct domains: art painting (A), cartoon (C), photo (P), and sketch (S).

For RGB datasets like GTSRB, CIFAR10, and CIFAR100, we adopt a pretrained VGG16 model in multi-domain training. In contrast, for grayscale datasets such as MNIST and FMNIST, lacking pretrained models, we design a convolutional network comprising three ResNet blocks, training it from scratch. And for DG, we adopt a pretrained ResNet18 for both datasets.
Initial codebook sizes are set to 32 for MNIST and 64 for the remaining datasets, with an equivalent number of codewords added in each subsequent iteration.
While additional iterations may converge within 5 rounds, we extend this to 20 for enhanced experimental clarity. The uncertainty evaluation is conducted 20 times using a dropout rate of 0.1, with thresholds $\gamma$ set at 0.3 for MNIST, 0.1 for FMNIST, GTSRB, and CIFAR100, and 0.2 for CIFAR10, to fine-tune performance. These experiments are performed on a machine with two NVIDIA A6000 GPUs.

Evaluation Metrics.
We calculate the mean Top-1 accuracy (mA) as across all silos for each distribution and all data to enable a straightforward comparison.
We evaluate entropy as model uncertainty as Eq. 7.
We also evaluate the perplexity (PPL) to show the utility of codewords as follows,

$$ $PPL=exp(-\sum_{i=1}^{N}p_{i}\log p_{i})$ (9) $$

where $N$ is the number of codewords, and $p_{i}$ denotes the probability of the $i$th codeword occurring.

Similar to mA, we evaluate mean entropy (mE) and mean perplexity (mP) across data silos.

### 4.1 Extensible Codebook

Discretization for Heterogeneous FL.
To show the effectness of discretization to tackle the data heterogeneity in FL, we design a toy experiment on MNIST.
Temporarily setting aside federated learning’s privacy considerations, we directly discretized the features for each client using the distinct codebooks based on its originating domain.
With this discretization of VQ-FedAvg, the mean accuracy was improved from 0.834 to 0.907 with the reduction of uncertainty, demonstrating the effectiveness of feature discretization in enhancing performance within a heterogeneous federated learning context, as shown in Fig. 3(a).

Figure: (a) Discretization works.
Refer to caption: extracted/5973906/code1.png

Extensible Codebook v.s. Static Large Codebook.
To validate our extensible codebook’s superiority over starting with a large codebook, we ensured both methods ended with the same number of codewords through experiments. Results on CIFAR100 showcased in Fig. 3(b) demonstrate the difficulties associated with a larger initial codebook in codeword selection for image features. Conversely, gradually expanding the codebook significantly improved codeword differentiation, 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 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 Fig. 3(c) with results from the MNIST dataset, showing enhancements across all metrics.

### 4.2 UEFL for Multi-Domain Learning

**Table 1: Comparison with different methods on heterogeneous data. mA denotes the mean of Accuracy for all data silos from the same distribution, and mE denotes the mean of Entropy. DisTrans lacks a Dropout layer, rendering it incapable of evaluating uncertainty.**
| Methods | Data | MNIST | FMNIST | GTSRB | CIFAR10 | CIFAR100 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| mA | mE | mA | mE | mA | mE | mA | mE | mA | mE |  |  |
| FedAvg | $\mathcal{D}_{1}$ | 0.874 | 0.212 | 0.801 | 0.246 | 0.670 | 0.623 | 0.676 | 0.172 | 0.110 | 1.74 |
| $\mathcal{D}_{2}$ | 0.848 | 0.231 | 0.825 | 0.232 | 0.677 | 0.634 | 0.622 | 0.178 | 0.072 | 1.86 |  |
| $\mathcal{D}_{3}$ | 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 |  |
| DisTrans | $\mathcal{D}_{1}$ | 0.856 | - | 0.721 | - | 0.898 | - | 0.721 | - | 0.289 | - |
| $\mathcal{D}_{2}$ | 0.799 | - | 0.705 | - | 0.900 | - | 0.719 | - | 0.261 | - |  |
| $\mathcal{D}_{3}$ | 0.789 | - | 0.694 | - | 0.897 | - | 0.659 | - | 0.251 | - |  |
| All | 0.815 | - | 0.707 | - | 0.898 | - | 0.699 | - | 0.267 | - |  |
| UEFL (Ours) | $\mathcal{D}_{1}$ | 0.951 | 0.120 | 0.857 | 0.147 | 0.95 | 0.0196 | 0.776 | 0.0192 | 0.362 | 0.728 |
| $\mathcal{D}_{2}$ | 0.885 | 0.196 | 0.848 | 0.188 | 0.964 | 0.0206 | 0.713 | 0.0245 | 0.335 | 0.624 |  |
| $\mathcal{D}_{3}$ | 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 |  |

We conducted comparative experiments on five datasets with introduced feature heterogeneity against leading algorithms, specifically the baseline Federated Averaging (FedAvg) (McMahan et al., ) and DisTrans (Yuan et al., ).
For accuracy comparison, DisTrans generally exhibits better performance than FedAvg, making it our primary point of comparison.
Regarding uncertainty comparison, because DisTrans lacks Dropout layers, precluding uncertainty evaluation, we exclusively compare uncertainty metrics with FedAvg.

Performance.
The results in Table 1 provide a comprehensive comparison, illustrating that our UEFL surpasses all other state-of-the-art (SOTA) methods in both accuracy and uncertainty reduction. Specifically, UEFL improves accuracy over FedAvg by 17.94% and DisTrans by 12.88% for the $\mathcal{D}_{3}$ distribution of the MNIST dataset.
And for uncertainty, our approach reduces uncertainty by 45.42% for the MNIST dataset’s $\mathcal{D}_{3}$ distribution.
Overall, our UEFL achieves accuracy improvements ranging from 3% to 22.1% over DisTrans.
Our UEFL improves uncertainty compared to FedAvg, achieving reductions by 38.83%-96.24%.
Figs. 4(a) and 4(b) details performance across individual data silos, highlighting our UEFL’s effectiveness in elevating the accuracy of last three silos and degrading the uncertainty.

Figure: (a) Accuracy comparison.
Refer to caption: extracted/5973906/silo1.png

Codewords Perplexity.
Fig. 4(c) presents a perplexity comparison between our UEFL and FedAvg, illustrating enhanced codebook utilization after assigning new codewords to $\mathcal{D}_{3}$. This adjustment not only benefits $\mathcal{D}_{3}$ but also improves the codebook utilization for $\mathcal{D}_{1}$ and $\mathcal{D}_{2}$.

Computation Overhead. Our approach introduces only a small codebook, thus incurring negligible memory and computational overheads. Specifically, for the CIFAR10 dataset, the parameter count for the baseline FedAvg model is 14.991M, whereas our UEFL model slightly increases to 15.491M, indicating a tiny memory increment of 3.34%. In terms of runtime, UEFL also exhibits a minimal increase from 16.154ms to 16.733ms (3.58% increase). These findings underscore UEFL’s suitability for deployment on edge devices.

### 4.3 UEFL for Domain Generalization

For domain generalization (DG) task, the trained model needs to be evaluated on an out-of-distribution domain and we follow the evaluation method in (Nguyen et al., ; Guo et al., ). Specifically, we perform “leave-one-domain-out” experiments, 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.

As shown in Table 2, our UEFL enhanced mean accuracy on the RotatedMNIST dataset, elevating it from 0.945 to 0.953. This performance exceeds that of FedSR (Nguyen et al., ) at 0.947 and FedIIR (Guo et al., ) at 0.95.
Similarly, on the PACS dataset, UEFL improved mean accuracy from 0.803 to 0.8453, surpassing FedSR’s 0.834 and FedIIR’s 0.837. These results underscore UEFL’s efficacy in tackling feature heterogeneity and superior performance on the federated domain generalization task, beating state-of-the-art methods.

Table: 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(McMahan et al., ), FedSR(Nguyen et al., ), FedIIR(Guo et al., )

### 4.4 Ablation Study

Imbalanced Clients.
We constructed an experimental setup with three data silos from $\mathcal{D}_{1}$ and one each from $\mathcal{D}_{2}$ and $\mathcal{D}_{3}$, totaling five silos. Our UEFL can also improve both accuracy (from 0.508 to 0.828) and uncertainty (from 0.256 to 0.105) in this scenario. Detailed results are in Appendix C.

Large Number of Clients.
We follow the settings in Guo et al. () to further segment the five training domains of Rotated MNIST into 50 sub-domains, each representing an individual client.
Our UEFL achieves the best mean accuracy of 0.9342, surpassing the performances of FedAvg at 0.908, FedSR at 0.912, and FedIIR at 0.93 as shown in Table 3, suggesting our UEFL is scalable for a larger number of clients.

Table: Table 3: Results show our UEFL is scalable for a large number of clients and beats all SOTA methods. We compare our approach with the following baselines: FedAvg(McMahan et al., ), FedSR(Nguyen et al., ), FedIIR(Guo et al., )

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 efficiency in K-means initialization. In Fig. 5(a) for GTSRB, initializing with 32 codewords provides comparable accuracy and uncertainty metrics.
For more complex datasets, we enhance selection capacity using codeword segmentation. Fig. 5(b) demonstrates that segmenting codewords into 4 parts leads to enhanced performance on CIFAR100.

Figure: (a) 32 codewords are sufficient.
Refer to caption: extracted/5973906/para1.png

Uncertainty Threshold.
In our UEFL, the uncertainty evaluator plays a pivotal role in identifying heterogeneous data without needing direct data access, with the threshold selection being critical. As illustrated in Fig. 5(c), a 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 increase computational
overhead.
Thus, in such cases, there is a trade-off between runtime and performance.

## 5 Conclusion

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-Based Extensible-Codebook Federated Learning (UEFL), leverages this extensible codebook through an iterative process that adeptly identifies data from unknown distributions via uncertainty evaluation and enriches the codebook with newly initialized codewords tailored to these distributions. The iterative nature of UEFL, coupled with efficient codeword initialization using K-means, ensures codewords are 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 adaptability. Our comprehensive evaluation across six prominent datasets showcases UEFL’s effectiveness, yielding accuracy enhancements ranging from 3% to 22.1% and significant reductions in uncertainty between 38.83% and 96.24%.

## Appendix A Our UEFL for Label Heterogeneity

Similar to Tang et al. (); Guo et al. (), we introduce label heterogeneity with dirichlet distribution ($\alpha=0.1$). The results in Table 4 show that our UEFL can also tackle the label heterogeneity when compared to FedAvg and performs better than VHL for CIFAR10 even if it cannot perform as well as FedBR.

**Table 4: Comparison with different methods on data with label heterogeneity.**
| Method | FMNIST | CIFAR10 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| FedAvg | VHL | UEFL (ours) | FedAvg | VHL | FedBR | UEFL (ours) |  |
| mA | 87.45 | 91.52 | 90.59 | 58.99 | 61.23 | 64.61 | 62.67 |

## Appendix B K-means Initialization

Fig. 6 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.

And to bolster the robustness of our methodology,
we dissect features into smaller segments—using
factors like 2 or 4—to pair them with multiple codewords,
thus covering the entirety of a feature vector as illustrated
in Figure 3. This segmentation exponentially increases the
codeword pool to $n^{2}$ or $n^{4}$, ensuring a robust representation
capacity.

Figure: Figure 6: (a) Kmeans initialization for heterogeneous data silos. (b) Workflow of discretizer.
Refer to caption: extracted/5973906/init2.png

## Appendix C Imbalanced Clients

As shown in Table 5, our UEFL also works for imbalanced data silos when there are three clients sampled from the same domain. Both accuracy and uncertainty get improved, especially for the third domain.

**Table 5: Our UEFL improves the performance of unbalanced data.**
| Data | Silo | FedAvg | UEFL |  |  |
| --- | --- | --- | --- | --- | --- |
| Acc | Entropy | Acc | Entropy |  |  |
| $\mathcal{D}_{1}$ | $s_{1_{a}}$ | 0.964 | 0.0312 | 0.952 | 0.0291 |
| $s_{1_{b}}$ | 0.936 | 0.0252 | 0.974 | 0.0170 |  |
| $s_{1_{c}}$ | 0.964 | 0.0499 | 0.944 | 0.0308 |  |
| $\mathcal{D}_{2}$ | $s_{2_{a}}$ | 0.796 | 0.1477 | 0.836 | 0.1261 |
| $\mathcal{D}_{3}$ | $s_{3_{a}}$ | 0.508 | 0.2560 | 0.828 | 0.1048 |

## Appendix D UEFL Optimization

Number of Codewords. We investigate the impact of varying
the number of initialized codewords within our extensible
codebook in Table 6, aiming to 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 or 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 for GTSRB and 32 for CIFAR100. Our findings indicate that, particularly 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.

**Table 6: Number of codewords. Experiment are on GTSRB dataset. ”-” denotes value close to 0.**
| #Codes | Data | $\mathbf{\mathcal{L}_{code}}\downarrow$ | mP$\uparrow$ | mE$\downarrow$ | mA$\uparrow$ |
| --- | --- | --- | --- | --- | --- |
| 8 | $\mathcal{D}_{1}$ | 3.61 | 4.78 | 2.02 | 0.257 |
| $\mathcal{D}_{2}$ | 3.68 | 4.86 | 2.01 | 0.265 |  |
| $\mathcal{D}_{3}$ | 3.38 | 4.81 | 2.03 | 0.249 |  |
| 16 | $\mathcal{D}_{1}$ | 3.41 | 10.46 | 1.36 | 0.515 |
| $\mathcal{D}_{2}$ | 3.55 | 10.54 | 1.37 | 0.506 |  |
| $\mathcal{D}_{3}$ | 3.23 | 10.61 | 1.39 | 0.486 |  |
| 32 | $\mathcal{D}_{1}$ | 0.127 | 25.91 | 0.0412 | 0.956 |
| $\mathcal{D}_{2}$ | 0.0885 | 25.42 | 0.0313 | 0.966 |  |
| $\mathcal{D}_{3}$ | 0.178 | 26.37 | 0.117 | 0.911 |  |
| 64 | $\mathcal{D}_{1}$ | 0.0975 | 26.79 | 0.0086 | 0.965 |
| $\mathcal{D}_{2}$ | 0.0853 | 26.32 | 0.0084 | 0.974 |  |
| $\mathcal{D}_{3}$ | 0.1907 | 27.23 | 0.0166 | 0.926 |  |
| 128 | $\mathcal{D}_{1}$ | 0.0512 | 38.73 | - | 0.954 |
| $\mathcal{D}_{2}$ | 0.0453 | 34.16 | - | 0.968 |  |
| $\mathcal{D}_{3}$ | 0.0726 | 43.42 | - | 0.917 |  |
| 256 | $\mathcal{D}_{1}$ | 0.0543 | 41.20 | 0.0043 | 0.962 |
| $\mathcal{D}_{2}$ | 0.0301 | 38.96 | 0.0054 | 0.959 |  |
| $\mathcal{D}_{3}$ | 0.0577 | 50.57 | 0.0103 | 0.904 |  |

Codebook Initialization. Section 3.1 highlights our 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 7 with
results from the MNIST dataset, showing enhancements
across all metrics.

**Table 7: The experiments were conducted on the MNIST dataset with 128 initialized codewords and segmentation factor 1. The model with K-means initialization outperforms without it.**
| Codebook | Data | $\mathbf{\mathcal{L}_{code}}$ | mP | mE | mA |
| --- | --- | --- | --- | --- | --- |
| w/o init | $\mathcal{D}_{1}$ | 0.6202 | 6.26 | 0.125 | 0.888 |
| $\mathcal{D}_{2}$ | 0.6063 | 5.99 | 0.296 | 0.554 |  |
| $\mathcal{D}_{3}$ | 0.6096 | 5.35 | 0.267 | 0.622 |  |
| w/ init | $\mathcal{D}_{1}$ | 0.0862 | 59.64 | 0.0935 | 0.945 |
| $\mathcal{D}_{2}$ | 0.0785 | 57.58 | 0.1604 | 0.906 |  |
| $\mathcal{D}_{3}$ | 0.0779 | 99.38 | 0.1509 | 0.929 |  |

Extensible Codebook v.s. Static Large Codebook. To
validate our extensible codebook’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 to 128 codewords and expanded twice, while
the static codebook was fixed at 512 codewords. Results
showcased in Table 8 demonstrate the difficulties associated
with a larger initial codebook in codeword selection
for image features. Conversely, gradually expanding the
codebook significantly improved codeword differentiation,
yielding better outcomes, such as 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 superiority.

**Table 8: Our extensible codebook (Extend) outperforms the static larger codebook (Static) on all evaluation metrics.**
| Codebook | Data | $\mathbf{\mathcal{L}_{code}}$ | mP | mE | mA |
| --- | --- | --- | --- | --- | --- |
| Static | $\mathcal{D}_{1}$ | 3.10 | 17.54 | 1.66 | 0.142 |
| $\mathcal{D}_{2}$ | 2.91 | 17.41 | 1.76 | 0.135 |  |
| $\mathcal{D}_{3}$ | 2.63 | 16.82 | 1.76 | 0.112 |  |
| Extend | $\mathcal{D}_{1}$ | 1.28 | 19.31 | 0.7822 | 0.375 |
| $\mathcal{D}_{2}$ | 0.976 | 27.42 | 0.6665 | 0.341 |  |
| $\mathcal{D}_{3}$ | 0.978 | 22.00 | 0.7112 | 0.304 |  |

Different Uncertainty Threshold.
In our UEFL, the uncertainty evaluator plays a pivotal role in identifying heterogeneous data without needing direct data access, with the threshold selection being critical. An optimal threshold enhances the model’s ability to distinguish between data silos, leading to quicker convergence. As illustrated in Fig. 7, a lower threshold imposes stricter criteria, pushing the model to achieve higher precision, thereby improving performance metrics. However, it’s important to recognize that beyond a certain point, further reducing the threshold may not significantly enhance outcomes but will increase computational overhead. Thus, in such cases, there is a trade-off between runtime and performance.

Figure: Figure 7: Results on GTSRB dataset with 64 initialized codewords with segment 1. Overall, a smaller threshold performs better.
Refer to caption: extracted/5973906/th2.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - The reported 11% accuracy for FedAvg on CIFAR-100 is unreasonably low, especially considering that the experiments used pretrained models. Could the authors clarify whether there were any specific preprocessing steps or hyperparameter choices that led to such low performance? Additionally, does UEFL still outperform FedAvg when both models are trained from scratch, without the benefit of pretraining?

- How does the growth of the codebook over multiple iterations affect the overall communication overhead in FL? Does the increased size of the codebook lead to significant additional communication costs during model aggregation, and how is this handled?

- How is the threshold for uncertainty (which triggers the expansion of the codebook) determined? Is this threshold set manually, or is it dynamically adjusted during training? Additionally, how sensitive is the method to this threshold?

- How does UEFL perform when the data distributions across silos are extremely divergent (for example, different modalities of data)? Does the codebook continue to expand indefinitely, or is there a mechanism to prevent overgrowth?
soundness: 2
strengths: - The idea of using a dynamic, uncertainty-driven codebook to handle data heterogeneity in federated learning is innovative. By extending the codebook based on uncertainty, the method adapts to diverse and unseen data distributions more efficiently than traditional methods.

- The approach of starting with a small, shared codebook and gradually expanding it based on data-specific uncertainties is efficient. This design helps in minimizing computational overhead while ensuring that the model can adapt to new distributions without retraining the entire model.
summary: The paper addresses the challenge of data heterogeneity in Federated Learning (FL) by proposing a novel method called Uncertainty-Based Extensible-Codebook Federated Learning (UEFL). The method introduces a dynamic codebook that maps latent features to discrete vectors, allowing for the extension of the codebook based on the uncertainty of data distributions across different data silos (clients). By iteratively evaluating the uncertainty of the global model's predictions, UEFL adds new codewords to the codebook to improve performance on previously unseen or highly divergent data distributions. The proposed method is tested on six datasets, showing significant improvements in both accuracy (3% to 22.1%) and uncertainty reduction (38.83% to 96.24%) over state-of-the-art approaches.
weaknesses: - In Table 1, the reported accuracy of FedAvg on CIFAR-100 is only 11%, which is surprisingly low. Typically, even with non-IID data, FedAvg achieves higher accuracy. This raises concerns about potential baseline suppression to exaggerate the performance of the proposed method. The authors need to clarify the experimental setup for FedAvg and ensure that the baseline methods are fairly represented. Without this clarification, the validity of the results is questionable.

- The experiments seem to rely on pretrained models. It would be essential to clarify whether the proposed UEFL method still performs well when trained from scratch. Relying on pretrained models can bias results, especially if the pretrained weights are already well-optimized for certain datasets. The authors should include experiments that train the models from scratch to validate the robustness and generalizability of UEFL. Additionally, the 11% accuracy for FedAvg on CIFAR-100 with a pretrained model is puzzling and further raises questions about the validity of the reported results.

- The paper primarily uses Monte Carlo Dropout for uncertainty estimation, which is a standard but limited approach. More advanced uncertainty quantification techniques (e.g., Deep Ensembles or Bayesian Neural Networks) could be explored or at least discussed as alternatives. The performance of UEFL might vary with different uncertainty estimation methods.

- The paper lacks a comprehensive analysis of the computational complexity of the proposed method. Although the authors claim that the codebook remains compact, it would be beneficial to quantify the computational and communication overhead in both the training and inference phases. A comparison of the complexity of UEFL with other FL methods in terms of memory usage and time complexity is missing.

- The paper mainly evaluates the method on six datasets, but does not provide a clear discussion of its scalability to larger numbers of clients. Given that FL is often applied to large-scale systems, it would be helpful to discuss how UEFL scales with increasing numbers of clients or more heterogeneous data distributions.

- Although the paper discusses the adaptive nature of the codebook, it would be useful to include an ablation study that explores the impact of different initial codebook sizes and how the codebook grows over iterations.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: It would be helpful if the authors could address my above questions.
soundness: 2
strengths: The authors address an important challenge in federated learning to tackle data heterogeneity and enhance model generalization across different distributions.

The presentation of the paper is easy to follow.
summary: This paper presents a method for federated learning, where the goal is to address data heterogeneity across clients and enhance generalization ability of federated learning model across different data distributions. Specifically, the method leverages a codebook of representations to improve accuracy and reduce prediction uncertainty in federated settings. The authors validate the proposed framework in multi-domain learning, and domain generalization in leave-one-domain-out experiments.

The proposed framework is as follows: for each client model, first encode the input data to latent features, extract code words via k-means clustering of these latent features, and then discretize the latent features using the nearest code word; predictions are subsequently made based on the coded (discretized) latent vectors. The framework also introduces an uncertainty-based adjustment mechanism, where clients with uncertainty above a certain threshold use a larger codebook with additional code words. The server model is updated as the average over all clients.
weaknesses: I have the following questions regarding the work:

1. My first question is about the rationale behind the framework. Could the authors elaborate on how and why discretizing latent representations through a codebook helps achieve the goals of improved generalization and uncertainty reduction?

2. Regarding the server model, could the authors clarify how the codebook of the server model is updated?

3. Is there any convergence guarantee of Algorithm 1? Specifically, does the framework ensure that all clients will reach an uncertainty level below the threshold? Could the authors clarify the mechanism described in lines 322–323?

4. Could the authors clarify why reducing uncertainty in prediction is a valid objective? If the model achieves high certainty on incorrect predictions, could this exacerbate challenges in improving model performance? Additionally, if reducing uncertainty is a primary goal, beyond the proposed method of expanding client codebook, have the authors considered training client models to the point of neural collapse [1]?

5. Distribution shift is also addressed in other studies, such as [2,3]. It would strengthen the evaluation if the authors could expand their comparisons to include more related methods.



[1] Papyan, Vardan, X. Y. Han, and David L. Donoho. "Prevalence of neural collapse during the terminal phase of deep learning training." Proceedings of the National Academy of Sciences 117.40 (2020): 24652-24663.

[2] Nguyen, A. Tuan, Philip Torr, and Ser Nam Lim. "Fedsr: A simple and effective domain generalization method for federated learning." Advances in Neural Information Processing Systems 35 (2022): 38831-38843.

[3] Zhang, Hao, et al. "FedCR: Personalized federated learning based on across-client common representation with conditional mutual information regularization." International Conference on Machine Learning. PMLR, 2023.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: It would be advantageous for the authors to include convergence curves relating to communication, computation, and the number of rounds. This would provide a clearer illustration of the effectiveness of UEFL.
soundness: 3
strengths: 1. Utilizing the codebook technique to address the feature shift in FL is innovative.

2. The authors conducted comprehensive experiments on multiple datasets, demonstrating the performance improvement of the proposed method in both traditional FL tasks and out-of-distribution generalization tasks.
summary: This paper proposes UEFL, which utilizes an extensible codebook to address feature shifts in FL. Numerical results demonstrate the superior performance of the proposed method in both multi-domain learning and out-of-distribution generation tasks.
weaknesses: 1. The algorithms presented in Section 3 do not align well with Algorithm 1. For instance, lines 226-228 are not thoroughly explained. It would be beneficial if the authors could provide additional details on the main pages.

2. The number of training epochs for the baseline models is set to be the same as that for UEFL. However, I am uncertain about the fairness of this approach, as (1) the optimal number of training epochs for UEFL may not be optimal for the baseline models; (2) the communication and computation costs associated with UEFL and the baseline models differ, which makes the comparison somewhat unfair.

Meta-review: Post-rebuttal, the majority of the reviewers still judged the contribution of the paper as fair. I appreciate the efforts the authors have put in their rebuttal, and the updates in the draft, very substantial and helpful in the appendices, but in the end the paper unfortunately does not pass the bar for acceptance. I can only encourage them further in updating the paper for resubmission, perhaps digging further in the theory that they so far cite from another paper.

Decision: Reject