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

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:

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.

Apologies for the confusion. Lines 226–228 (in the reviewed manuscript) explain the process of codebook augmentation in each iteration. Specifically, we identify data from heterogeneous distributions by evaluating uncertainty against a pre-established threshold. When such data are detected, the codebook is expanded with $v$ new codewords, initialized using K-means. These new codewords are exclusively accessible to the corresponding heterogeneous clients, increasing the codebook size for the $k$th client from $v_k$ to $v_k + v$.

We have updated the notation for the number of codewords to $v$ and included additional details in the revised manuscript for improved clarity (Lines 206–210).

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

We selected the same number of training epochs as UEFL because UEFL requires 1 or 2 additional iterations, whereas other baselines converge earlier. The numbers reported in our tables represent the optimal performance achieved within these training epochs. For a fairer comparison, we also conducted all experiments under the official experimental setup, excluding the additional epochs from UEFL. The results, presented in the format Accuracy (Uncertainty), are as follows:

Owing to the stochastic process, there are slight differences compared to the previous results. However, the overall performance remains consistent, and the key finding remains unchanged: UEFL outperforms all baselines across all datasets. These results have been included in the appendix of our revised manuscript. (Lines 752-762)

Regarding communication and computation costs, our approach introduces only a small codebook, resulting in negligible memory and computational overhead. The detailed comparison with FedAvg is presented in the following table. For UEFL, we report the final model size with 256 codewords, which is the largest number across all datasets in our experiments. Specifically, the parameter count for the baseline FedAvg model is 14.991M, while UEFL increases slightly to 15.491M, reflecting a modest memory increment of just 3.34%. Similarly, the runtime shows a minimal increase from 16.154ms to 16.733ms, representing only a 3.58% increase.

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

Thank you for your suggestion. We have incorporated the convergence curves (, ) for various metrics (including accuracy and uncertainty) on the MNIST, FMNIST, GTSRB, and CIFAR10 datasets into the revised manuscript. These curves clearly highlight the improvements brought by UEFL and effectively demonstrate its advantages. (Lines 825–882)

Regarding communication and computation costs, we have also included the aforementioned table in the revised manuscript. (Lines 427-431 and Lines 810-822)

Thank you for your thoughtful feedback, including your recognition of the motivation and presentation of our work. We also greatly appreciate your insightful critique regarding the limitations of our study. In response, we address these concerns as follows:

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?

The rationale mainly comes from this previous work [1]. Theoretically, the inclusion of the discretization process offers two key advantages: (1) enhanced noise robustness, and (2) reduced underlying dimensionality. These benefits are demonstrated in the following two theorems.

Notation: $\boldsymbol{h}$ is input vector, $\boldsymbol{h}\in\mathcal{H}\in\mathcal{R}^m$. $L$ is the size of codebook, $G$ is the number of segments, $q(\cdot)$ is discretization process, $\phi(\cdot)$ is any function (model). Given any family of sets $S=\\{S_1, ..., S_K\\}$ with $S_1, ..., S_K\subseteq\mathcal{H}$, we define $\phi_k^S$ by $\phi_k^S= \mathbb{1}\\{\boldsymbol{h}\in S_k\\}\phi(\boldsymbol{h})$ for all $k\in[K]$, where $[K]=\\{1, ..., K\\}$. And we denote by $(Q_k)_{k\in[L^G]}$ all the codewords.

Theorem 1: (with discretization) Let $S_k=\\{Q_k\\}$ for all $k\in[L^G]$. Then for any $\delta>0$, with probability at least $1-\delta$ over an iid draw of $n$ examples $(\boldsymbol{h}\_i)_{i=1}^n$, the following holds for any $\phi:\mathcal{R}^m\rightarrow\mathcal{R}$ and all $k\in[L^G]:$ $if$ $|\phi_k^S(\boldsymbol{h})|\leq\alpha$ for all $\boldsymbol{h}\in\mathcal{H}$, then

$

\left\|\mathbb{E}_h[\phi_k^S(q(h, L, G))]-\frac{1}{n} \sum^n\_{i=1}\phi_k^S(q(h_i, L, G))\right\|=\mathcal{O}(\alpha\sqrt{\frac{G\text{ln}(L)+\text{ln}(2/\delta)}{2n}}),

$

where no constant is hidden in $\mathcal{O}$.

Theorem 2: (without discretization) Assume that $||\boldsymbol{h}||\_2 \leq R\_{\mathcal{H}}$ for all $\boldsymbol{h}\in\mathcal{H}\in\mathcal{R}^m$. $Fix$ $\mathcal{C}\in argmin_{\overline{\mathcal{C}}}\{|\overline{\mathcal{C}}|:\overline{C}\subseteq\mathcal{R}^m, \mathcal{H} \subseteq \cup_{c\in\overline{\mathcal{C}}}\mathcal{B}[c]\}$ where $\mathcal{B}[\boldsymbol{c}]=\{ \boldsymbol{x}\in\mathcal{R}^m: ||\boldsymbol{x}-\boldsymbol{c}||\_2\leq R_{\mathcal{H}}/(2\sqrt{n}) \}$. Let $S_k=\mathcal{B}[\boldsymbol{c}_k]$ for all $k\in [|\mathcal{C}|]$ where $\boldsymbol{c}\_k\in \mathcal{C}$ and $\cup_k\{\boldsymbol{c}\_k\}=\mathcal{C}$. Then, for any $\delta >0$, with probability at least $1-\delta$ over an iid draw of $n$ examples $(\boldsymbol{h}\_i)\_{i=1}^n$,the following holds for any $\phi : \mathcal{R}^m \rightarrow \mathcal{R}$ and all $k\in[|\mathcal{C}|]:$ if $|\phi_k^S(\boldsymbol{h})|\leq\alpha$ for all $\boldsymbol{h}\in\mathcal{H}$ and $|\phi_k^S(\boldsymbol{h})-\phi_k^S(\boldsymbol{h}')|\leq\varsigma_k||\boldsymbol{h}-\boldsymbol{h}'||\_2$ for all $\boldsymbol{h}, \boldsymbol{h}'\in S_k$, for all $\boldsymbol{h}, \boldsymbol{h}'\in S_k$, then

$

 \left\|\mathbb{E}\_{\boldsymbol{h}}[\phi_k^S(\boldsymbol{h})]-\frac{1}{n} \sum\_{i=1}^n \phi_k^S(\boldsymbol{h}_i) \right\|=\mathcal{O}(\alpha\sqrt{\frac{m\text{ln}(4\sqrt{nm})+\text{ln}(2/\delta)}{2n}}+\frac{\overline{\varsigma_k}R\_{\mathcal{H}}}{\sqrt{n}}),

$

where no constant is hidden in $\mathcal{O}$ and $\overline{\varsigma_k}=\varsigma_k(\frac{1}{n}\sum_{i=1}^n\mathbb{1}\{h_i\in \mathcal{B}[c_k]\})$.

Based on these two theorems, we can determine that the performance gap between training and test data is smaller when discretization is applied, due to the following two points:

• There is an additional error without discretization (i.e. $\frac{\overline{\varsigma_k}R_{\mathcal{H}}}{\sqrt{n}}$) in the bound of Theorem 2. This error disappears with discretization in the bound of Theorem 1 as the discretization process reduces the sensitivity to noise.
• The discretization process reduces the underlying dimensionality of $m\text{ln}(4\sqrt{nm})$ without discretization (in Theorem 2) to that of $G\text{ln}(L)$ with discretization (in Theorem 1). Since the number of discretization heads $G$ (eg. $G$ is 1, 2, or 4 in our case) is always much smaller than the number of dimensions $m$, the inequality $G\text{ln}(L)\leq m\text{ln}(4\sqrt{nm})$ consistently holds.

This is consistent with the experimental results presented in Section 4.1 of our manuscript. Specifically, with the discretization process, the mean accuracy improved from 0.834 to 0.907, accompanied by a reduction in uncertainty (Fig. 3a). The details are included in Lines 671-709 of the revised manuscript.

[1] Liu, Dianbo, et al. "Discrete-valued neural communication." Advances in Neural Information Processing Systems 34 (2021): 2109-2121.

We appreciate the recognition of the novelty and efficiency of our method. Additionally, we value the insightful critique regarding the limitations of our work. In response, we address these issues below:

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 performance heavily depends on the experimental setup, particularly data heterogeneity. In our manuscript, for the CIFAR-100 dataset, we adopted the same setup as other datasets, introducing high data heterogeneity by rotating the images with large angles. To ensure fairness, we also systematically conducted experiments under an alternative setup with lower data heterogeneity.

To provide a detailed comparison with the baseline, we utilized a VGG16 backbone to test CIFAR-100 across multiple settings:(1) Local Training: all data trained together without a distributed setting; (2) FedAvg (w/o hete): CIFAR100 split into 5 clients, each with 10,000 images, to evaluate FedAvg performance; (3) FedAvg (w/ hete): images for the 5 clients were rotated by -30°, -15°, 0°, 15°, and 30°, respectively, to introduce data heterogeneity, and FedAvg performance was evaluated; (4) UEFL (w/ hete): tested under the same heterogeneous setup. By presenting results across these progressively changing setups, from local training to heterogeneous federated learning, we demonstrate the validity of our experimental approach. The results are as follows:

The accuracy of both local training and federated learning aligns with publicly reported numbers. Under lower data heterogeneity, both FedAvg and UEFL achieve higher accuracy, with UEFL consistently outperforming FedAvg in this experimental setup. The results are updated in Lines 714-730.

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.

For this point:

• As mentioned in our manuscript, for grayscale datasets such as MNIST and FMNIST, where pretrained models are unavailable, we design a convolutional network consisting of three ResNet blocks and train it from scratch. This clarification has been highlighted in blue in the revised manuscript (Lines 335–337).

• Additionally, we conducted experiments on CIFAR-100 under the aforementioned experimental setup, training from scratch without using pretrained weights. The detailed results are presented in the table above. UEFL continues to outperform the baseline when training from scratch.