Paper: FLAT-Chat: A Word Recovery Attack on Federated Language Model Training (ICLR-2024)

Reviewer comment (Writing):
**The description of the proposed method can be hard to read at times:**  
I know a lot about this particular area of research and I still struggled to follow the presentation of Section 3 (the technical contribution section). To this end, in my opinion, the paper will really benefit from a paragraph (probably coupled with a summary figure) that summarizes the steps of the proposed method early on in Section 3, so that it is easier to follow what the paper is trying to achieve through the different subsections of Section 3. It needs not be long, consider something like the beginning of my paper summary above. Similarly, presenting the full algorithm at the end of Section 3 will help a lot in understanding how the different pieces of the algorithm fit together. Further, the paper will also benefit from giving more intuitive explanations of its steps throughout. One example of this will be to present Eq. 9 before Theorem 3 to make it intuitively clear where the GMM pieces come from in Eq.6. Finally, there are several key missing from Section 3. It should explicitly state that the GMMs and the regression model on $\|\mathcal{T}\|$ need to be fitted on auxiliary data, and how estimating the number of unique tokens from the regression model is used together with the ranking to provide the set of recovered tokens. It should also state that during the LLM training, when multiple tokens are predicted, the CE loss is summed across all of them which mathematically is equivalent to the label recovery from a large batch.

Author response:
We sincerely appreciate these suggestions regarding rearrangements and supplements. We will incorporate them in our revision.

Reviewer comment (Novelty):
**Citating and comparing to prior work:**  
The paper should cite and compare against prior label reconstruction attacks outside of RLG [5]. In particular, [1] can be used for recovering the set of unique tokens, while [2,3] can be used to recover the counts as well. Comparing against [1-3] is absolutely crucial, in my opinion, for accepting this paper, as those methods would work fast for large vocabularies and long sentences, unlike RLG, and have been shown to be effective at recovering labels to very good accuracy. Further, [2], in particular, is very closely related to FLAT-CHAT, as it derives the same "flattening" operation the authors claim as a contribution in the text. To this end, the authors should not claim the flattening operation as a contribution and instead clearly mark the derivation presented there as equivalent to the one made in [2].  
Given the similarities to prior work, the authors should also consider including an explicit discussion of how their method differs from prior work. Finally, the authors acknowledge that FILM [4] can be applied to the same problem the authors consider but from the input side of the network. Yet, they do not provide a comparison. While beating it is not required for acceptance (due to the different requirements the attacker has), comparing against it is a good idea.

Author response:
Thank you for pointing out the literature.

[1] and [3] are for image reconstruction, focusing on image data in a continuous space while our work focuses on text data in a discrete space.

[2] shares similar insights to our Lemmas 1 and 2. We have compared this baseline (coined Abs.) in Table 4. We will acknowledge our Abs. is equivalent to the suggested work. We would like to highlight our progress compared with [2]: (1) we have extended the theory and method (see GMM score and its results) and (2) our approach manages to estimate the number of tokens used in the recent training batch, which is often omitted in many related works.

[4]  FILM targets sequence recovery while our work focuses on *industrial-scale* word recovery attack. Nonetheless, our work could serve as the first *Bag-of-Words* Extraction step of FILM.

Reviewer comment (Evaluation):
**The attack setting:**   
Label leakage attacks like [2] and [3], are capable not only of recovering the set of unique tokens in input data but also their counts. The authors should provide a discussion on whether counts are important from LLM privacy point of view.   
Further, the authors should better motivate their attacker's goal in general. While privacy is indeed violated by knowing the set of tokens fed to the network from a purely theoretical point of view, I would reasonably think that a large percent of the vocabulary tokens occur in a large batch of long excerpts of text anyway, and when the recovery has a precision of 0.85 and recall of 0.5 it will be very hard from a practical perspective to gain any reasonable sensitive information. That is, I expect the rank of rare tokens, which tend to be more private, to be lower in your method due to their lower occurrence rates. I also expect that the recall will be much lower than 0.5 for labels that are in the middle of the ranking.  Thus, in such a situation, the attacker will obtain that words like "the", "I", "you" are present in the batch with high accuracy, but will rarely obtain, let's say, a phone number. The problem gets even worse when considering the fact that LLMs are trained on tokens and not full words.

Author response:
We argue recovering the appearance of tokens is the first and the most essential step for information leakage. The frequency of a token is arguably less important than know it is used or not. Nonetheless, we appreciate the suggestion and would like to leave it to future work.

Reviewer comment (Evaluation):
**Bad evaluation results:**  
The results shown in the experiment do not convince me in the superiority of the proposed method. In particular, RLG consistently and by big margins results in better reconstructions than FLAT-CHAT if RLG is in the mode where it is applicable (\|\mathcal{T}\| < D). [1-3], which do not have such restrictions and tend to work much faster than RLG, might, therefore, turn out to be much better than FLAT-CHAT.  
Even outside of these concerns, I find the precision of 0.85 and the recall of 0.5 in Table 2 and the 0.7 precision and 0.85 recall numbers in Table 3 not that convincing in terms of their practical attack relevance as outlined above.

Author response:
Actually, the out-of-scope size of training data is quite common, imagining industrial FL scenarios using large training batches. DLG is not aware of the potential issue of a numerical value exceeding the range of $D$ which is usually several hundred. Our estimation model based on the parameters from GMM could solve the issue very well, according to Figure 3.

Reviewer comment (Theory):
**Suggestions — Modeling probability distributions:**  
Essentially, the method proposes to model the $p_{i,j}$ as Gaussian distribution, which as $\|\mathcal{B}\|->\infty$ get closer to the truth but since $0\leq p_{i,j}\leq 1$ is a probability the approximation for finite $\|\mathcal{B}\|$ is very bad. This is also reflected in the authors' shown negative clusters in the figures of Appendix A. The authors can consider modeling the $\log p_{i,j} $ as Gaussian (but $\alpha_i$ still as Gaussian). In my quick tests, this reflected the shown negative cluster pdf shapes much better.

Author response:
Thanks for the suggestion, this method may serve as a better estimator. We will consider adding this to our future work.

Reviewer comment (Theory):
**Suggestions — Ranking function improvement:**  
The authors propose to use Equation 12 as a ranking function. If a proper prior $p(s_t|t \in \mathcal{B})$ is used, Equation 12 can be used as a decision criterion instead, which will eliminate the need for using regression to fit $\|\mathcal{T}\|$. This can possibly improve the performance of the method.

Author response:
There is an issue in the suggested method to estimate $|T|$: the $|B|$ is the suggested formula, $\frac{|B|}{(|V|-1)|B|}$, will be canceled. We did not find promising methods in predicting $|T|$ in [1-3], especially under the setting of large batch size and large vocabulary size.

Reviewer comment (Theory):
**Nits — Equation 13 constants:**   
In the first part of Eq. 13, $\sigma_n$ and $\sigma_p$ should be switched in the normalization constants of the Gaussians.

Author response:
We sincerely appreciate your suggestions, and we will incorporate them into our revision.

Reviewer comment (Theory):
**Nits — Consistency in aggregation assumptions:**   
Equations 6 and 9 assume sum instead of mean gradient aggregation. Equation 7 assumes a mean instead of sum. This needs to be made consistent throughout the paper.

Author response:
We sincerely appreciate your suggestions, and we will incorporate them into our revision.

Reviewer comment (Evaluation):
**Attack significance is low** because the order of tokens cannot be recovered. As the attack relies on gradient distribution of positive and negative tokens, the tokens' order information is hidden and not recovered. Although word distribution can leak partial privacy, in my opinion, this information is important to infer privacy underlying the training text. As a simple example, the two texts X = "A is good, B is bad", and Y = "A is bad, B is good" have the same word distribution but totally different meaning. I suggest the author to focus on or highlight specific scenarios where the word distribution can leak sufficient privacy. For example, it is possible to conduct an end-to-end case study showing how recovered tokens can lead to a more severe consequence.

Author response:
Revealing the training word types could also infringe serious risk for the scenarios requiring brief context, like healthcare diagnosis or personal identifiable information (PII), including credit card numbers and social security numbers. The significance of our work lies in that we are the first work that scaled up the word reconstruction attack to *industrial* settings, such as more than 10 K tokens in a batch, vocabulary size larger than 50 K and extremely fast inference speed. Furthermore, consolidated word reconstruction could serve as the first step of sequence reconstruction attacks, followed by iteratively swapping tokens and training them within a language model, as adapted in the prior work FILM and LAMP [1].

Reviewer comment (Experiments):
**The technical challenge is not clear.** As Fig.2 shows, most negative tokens have vector $s$ value between 0 and 0.02. While I appreciate the authors' efforts in visualization, it makes me doubt whether the GMM is necessary. As the next word prediction resembles to classification, a naive baseline can apply iDLG-similar approach to directly identify (for example, with threshold) the trained words (positive tokens). Note that iDLG also leverages the last layer's gradient to infer the labels of trained samples. In this sense, the GMM is only used to better classify positives and negatives. I suggest the authors to make the attack motivation and challenges more clear in the paper.

Author response:
The focus of the iDLG is to infer the input data X, and it requires gradients from all layers. The usage of last layer gradients is just to infer the classification label for better gradient matching. However, our method focuses on inferring the output data Y from a single last layer's gradients.

Reviewer comment (Novelty):
**Problem importance is unclear.** From the main text, I cannot see that FL is a common solution for training/finetuning LMs, especially the large ones. Although the authors have provided a long list of related works of training data inference attack in FL, I think it is still important to show that FL is or will be applied by organizations through real-world examples or case studies. The only application I can imagine is using FL on mobile keyboard to predict the user's input behavior more accurately, but I'm not sure whether it trains such LMs. According to my experience, finetuning current LMs requires relative large memory, which is impractical to proceed on edge devices. Please illustrate potential FL applications for LM training.

Author response:
The LM training can be under the scenarios like training LM for medical reports, financial reports, keyboard as mentioned etc. In many scenarios, the text corpus would be kept in clients for privacy concerns and FL will be their first choice in training shared powerful LMs (which could be used for many purposes, such as classification, QA, NLI, NLG, etc.).

Reviewer comment (Experiments):
**Comparison with more baselines is needed.** I note that in Table 1 a recent work FILM also infers trained words but is not compared in Section 4.2.1. I also notice that there is a slight difference between FILM and this work in terms of $\Delta W$ but I think under FL setting the FILM can also work. Please consider compare with this attack or clarify why it is not suitable for comparison.

Author response:
FILM targets sequence recovery while our work focuses on *industrial-scale* word recovery attack. Nonetheless, our work could serve as the first *Bag-of-Words* Extraction step of FILM.

Reviewer comment (Evaluation):
**Defense (DP-SGD) can mitigate the attack, further reducing the attack significance.** To be honest, I'm quite surprised that small noises added by DP-SGD can mitigate the attack, which is different from the conclusion in (Gupta et al. 2022). This means that this previous attack is more powerful than proposed attack because DP-SGD can not defend it without degrading the model utility.

Author response:
The significance of our work is (1) We propose a new word recovery attack which works on *industrial-scale*, with theoretical proof and extensive empirical evidence; (2) We demonstrate the risk could be mitigated by DP-SGD. Our work enhances the motivation of using a proper security schema in FL. It is unfair to discredit a novel attack merely because it could be potentially defended, e.g., the many attack methods, FILM and LAMP [1], proposed by recent literature could also be defended by DP-SGD to some extent.

Reviewer comment (Reproducibility):
What is the learning rate used in attack and DP-SGD? What is the resultant budget ($\epsilon$, $\delta$)?

Author response:
Our experiments use a learning rate of $5 \times 10^{-4}$, as detailed in the footnote of Section 4.1, which introduces the experimental setup.  We will include the privacy budget in our revision.

Reviewer comment (Presentation):
What does the 'Loss' in Figure 4 mean? Training loss or validation loss?

Author response:
The ‘Loss’ in Figure 4 denotes the averaged cross-entropy loss, indicating the utility (how efficiently it operates in a language prediction task) of a language model. We will include an explanation in our revision.

Reviewer comment (Writing):
Some writings can be simplified, e.g. the lemmas and their proofs in Section 3.2 are simple rearrangement using some basic linear algebra which can be condensed in Equations and will not impact their readability. The theorems and the body texts are interleaved which makes the explanation of the attack less easy to follow.

Author response:
The proof of the properties is one of the key findings in this work as they ground the success of our word recovery attack. We sincerely appreciate your suggestions, and we will rearrange this content in our revision for better readability.

Reviewer comment (Experiments):
In common practice, when training language models, the parameters of the embedding layer and the last layer are typically shared, i.e. they have the same gradients. It would be a stronger attack if this more common scenario is considered.

Author response:
Yes, we have conducted tests on this setting in our preliminary study on MT tasks and our method has demonstrated successful outcomes. However, it is important to note that the share-embedding setting exposes both embedding and the last linear layers to data reconstruction attacks.

Reviewer comment (Evaluation):
The attack is limited to inferring the bag of words while the order of the words cannot be recovered.

Author response:
Word reconstruction could also pose a privacy risk in instances like exposing trade secrets, healthcare patient diagnosis, or personal identifiable information (PII) such as credit card numbers. Furthermore, consolidated word reconstruction could serve as the first step of sequence reconstruction techniques, followed by iteratively swapping tokens and training them within a language model, as adapted in the prior work FILM and LAMP.

Reviewer comment (Experiments):
How would tying the weights between input embedding and output layer change the performance of the attack?

Author response:
Our implementation is in the context of weight tying and has demonstrated successful outcomes. We will add such an explanation to our revision.

Reviewer comment (Experiments):
Another potential defense is secure aggregation, where the server can only observe the aggregated gradients instead of individual’s. How might this impact the attack? Could the adversary still infer useful information when the set of participants is large?

Author response:
Thanks for your question. That is why we need to consider the setting with more training tokens involved. It is worth noting that aggregating (say averaging) gradients from multiple batches (say from various clients) is equivalent to the gradients acquired from the `aggregated' (say concatenating all samples in these batches as a single batch) larger batch.

Reviewer comment (Experiments):
Experimental results could be more comprehensive. In particular, more exploration (e.g. of larger batch sizes) would establish the failure mode of the approach.

Author response:
(1) We have covered some experiments with the *failure* performance of RLG, such as 1,500 and 2,000 tokens in Table 2. (2) When scaling the batch size, the attacker will still be valid as we find **Prec.** score of FlatChat is kept stable or even grows in some cases (see Tables 2 and 3), which benefits from a relatively conservative batch size predictor.

Reviewer comment (Experiments):
I am also curious about the gaussian mixture model of the word types; Does the frequency of each word in the batch impact the quality of the fit? Experiments that demonstrate robustness in this scenario would be helpful in establishing generality of the approach.

Author response:
We consider it as a question out of curiosity instead of an argument of weakness. It is implicitly encoded in our Theorem 3, as scores of high-frequency words will be at the `centre’ of the positive cluster.

Reviewer comment (Experiments):
Further experiments which show performance of freeze/dp-sgd for language modelling would also help contextualize the benefits and drawbacks of the proposed defenses.

Author response:
Yes, we have demonstrated the performance of both freeze and dp-sgd on Machine Translation task in Figure 4.

Reviewer comment (Experiments):
I am a bit confused why results for Scratch with the task of Large Language Modelling are not included.

Author response:
In many LLM application scenarios, we directly use well-trained LLM instead of LLM from scratch. Additionally, attacking Scratch is a relatively easier experimental setting for attackers as demonstrated in Machine Translation tasks.

Reviewer comment (Evaluation):
Is the gaussian mixture model accurate at every epoch of fine-tuning? Or is it only at the first epoch?

Author response:
GMM could be easily generalized to various settings:
1. The feasibility of the Gaussian Mixture Model is proved by our Theorem 3 based on Theorem 4 (CLT) 
2. FlatChat empirically works on well-finetuned models, such as MT’s FINETUNE and GPT-2. 
3. Because the data leakage in the later epochs is of less interest given the data has already been disclosed in the earlier epochs, we focus on the token recovery attack in the early training epoch. 
4. Any gradient-based reconstruction attacks will fail in an extreme case, say the model overfits a batch of data leading to zero loss and zero gradient vectors.