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:

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.

We 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.

Furthermore, 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.

Nonetheless, 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]. We will leave the detailed exploration and integration of such privacy-preserving techniques for further study.

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.

We 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.

Although 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.

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.

While 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. In the revised supplementary materials, we will include full derivations based on Hoeffding Inequality for Theorem 1 and 2.

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.

Thank 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

Thanks once again for the constructive comments and valuable suggestions, which have significantly enhanced the quality and clarity of our manuscript.

[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.

[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.

We appreciate the insightful comments and suggestions. In response, we address these issues below:

Scalability to thousands of clients is not tested. Estimate overhead w.r.t. the number of clients

While 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:

$

Overhead(i) = \frac{i\times 0.125}{14.991}\times 100\\%

$

UEFL typically needs 1-3 iterations and we set the maximum to 5 (Section 3.2), so the introduced overhead is at most 4.17%.

In some FL settings, stricter constraints (e.g., differential privacy, homomorphic encryption) exist

Thanks 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.

The theoretical section ... rely on IID assumptions, while real FL data is often non-IID

In 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

With discretization:

$

\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}}),

$

Without discretization:

$

\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}}),

$

Here, $\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$. Therefore, 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.

Figure 9 ... no principled way to set threshold

In 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.

Not analyze long-term codebook growth ... analysis of the size of the codebook on the uncertainty

Computation is discussed above.

Codeword 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.

The plots are blurry. Use vector illustration

Thanks. 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

The running title needs an update.

We have updated it to: "UEFL: Uncertainty-Based Extensible Codebook Federated Learning".

The experimental setting is unclear ... Use both GPUs or only one?

All experiments utilized only one GPU, added to the revision.

How does performance degrade if the codebook is not expanded enough?

Then, the model lacks representational capacity, resulting in reduced accuracy and higher uncertainty (Figure 6). However, UEFL typically converged after 1-3 iterations.

Incorporate newer techniques like FedPer, FedRod, FedBABU, or FedDyn

Here is the comparison with FedRod & FedDyn:

UEFL outperforms them. More results will be included in the updated version.

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?

While the encoder is strong, the classifier is still vulnerable to domain shifts. UEFL improves performance via discretization, as shown in Table 1.

[1] Huh et al., Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks, ICML 2023.

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:

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?

In 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.

Codebooks 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.

We will clearly clarify these points in the revised manuscript to explicitly address potential privacy concerns. We deeply appreciate the thoughtful review and suggestions.

[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.

[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.

Thanks for the insightful comments and suggestions. We address your concerns as follows:

While the theorems ... connection to UEFL’s empirical success is not explicitly discussed

We provide a detailed analysis in response to the later theoretical questions. Please refer to that.

The paper mainly uses rotation transformation ... not cover more complex distribution drift situations

Other forms of heterogeneity are also included:

• Feature space shifts: Tables 2 and 3 (e.g., PACS domains)
• Uneven class distributions: Table 12 (Label heterogeneity)

Lack of detailed discussion on the impact of different codebook sizes on model stability

We 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).

Only used 5 sub models ... not explore whether increasing the number of sub models would affect the stability of the assessment

We extended the ensemble size to 20 and observed the following on MNIST:

Appendix B ... not provide a detailed generalization error curve

We have plotted the curve: https://ibb.co/xymvPDT

Appendix F ... does not explore the impact of communication costs

In 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.

More extensive experiments are still needed to verify their applicability

We’re extending this to medical data and other domains in future work.

Essential References Not Discussed

Thanks. Specifically for FedProx, the results are as follows:

The paper provides a theoretical analysis ... not consider the Non IID distribution in FL

In non-IID FL, then

With discretization:

$

\mathcal{O}( \alpha \sqrt{ \frac{G \ln L + \ln(2K/\delta)}{2n}} + \frac{\nu^{(q)}}{\sqrt{n}}),

$

Without discretization:

$

\mathcal{O}( \alpha \sqrt{ \frac{m \ln(4\sqrt{nm}) + \ln(2K/\delta)}{2n} } + \frac{\overline{\varsigma} R\_\mathcal{H} + \nu}{\sqrt{n}}),

$

Here, the KL divergence $\nu^{(q)} < \nu$. Therefore, discretization explicitly mitigates the effects of data heterogeneity.

To conserve space, more details are included in our response to Reviewer idfz.

Not provide a mathematical analysis of the effect of codebook size on convergence

The 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. This improves codeword utilization, reduces quantization error, and accelerates convergence (Fig. 3(b)).

Limited exploration of Label Heterogeneity

We further test different dirichlet distributions as follows,

Only uses K-means for initialization ... not explore other possible initialization methods

Thanks. We agree that PCA or contrastive-based initialization is promising and will explore them in future work.

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?

No. 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}$. With 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).

Can Theorem 1 be extended to account for dynamic codebook growth?

Yes. 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.

Experiment conducted under more extreme Non IID settings (α=0.01)?

The results for Dirichlet distribution α=0.01 are discussed above.