DAGER: Exact Gradient Inversion for Large Language Models

\newcommand{\Rj}{\textcolor{green}{ju5t}}$$\newcommand{\RL}{\textcolor{red}{L4TG}}We would like to thank reviewer $\Rj$ for the positive review and the thorough and insightful questions. We are happy they found our work effective, and the empirical results impressive, highlighting the significant improvement over state-of-the-art gradient leakage attacks. We have compiled a list of answers to any remaining concerns below.

Q1. What are the technical assumptions for DAGER? Is there any requirement on the loss function, i.e. MSE or cross-entropy? Are there any requirements on the differentiability?

We list the technical assumptions regarding the assumptions in Q2 of the general response. In short, we list the underlying assumptions and explain that they are much weaker compared to prior work. We want to emphasize that DAGER’s essential assumption is that $\frac{\partial \mathcal{L}}{\partial\bm{Q}_l}$ is full-rank. As such, the (sub-differentiable) loss function needs to non-trivially depend on all inputs, which is satisfied for any reasonable one.

The full-rankness of $\frac{\partial \mathcal{L}}{\partial\bm{Q}_l}$ is also the limiting factor for any activation function. As they are applied only after the MLP at the end of a transformer block, they are less likely to directly affect $\frac{\partial \mathcal{L}}{\partial\bm{Q}_l}$. Further, Dimitrov et al. [A] have empirically demonstrated that the induced sparsity from a ReLU activation does not induce low-rankness even on the weight directly preceding it. Therefore, we reaffirm that DAGER can be applied to any reasonable activation, and that we do not depend on any particular properties.

Finally, as demonstrated in Q2 of the general response, we can indeed handle MSE loss or ReLU activations.

Q2. Why does DAGER require b<d? Is this practical given common embedding space sizes of 512 to 1024? Dagger requires $b<d$ to ensure the low rankness of $\frac{\partial\mathcal{L}}{\partial\bm{W}^Q_l}$ required for our approach. However, despite this limitation, we significantly outperform every prior work by a wide margin, both in batch size and in sequence length on well-known datasets. Further, we want to point out that the embedding dimensionality for GPT-2 and LLaMa-2 (7B) are already 768 and 4096 respectively, (See Table 4 in the Appendix), with newer models often having even larger embedding spaces, e.g. LLaMa-3.1 (405B) has $d = 16,384$. Especially for these larger embedding dimensions we believe it to be quite rare to encounter this limitation in current federated learning applications. For example, in the FedLEGAL benchmark [A] with inputs of similar length to the Rotten Tomatoes dataset a batch size of up to $B=64$ poses no issue.

Q3. How do linearly dependent token embeddings affect DAGER?

It is clear that the token embeddings in a vocabulary are linearly dependent, as the vocabulary size usually exceeds the embedding size. However, this is not a concern for DAGER, as we are only interested in the embeddings present in the batch we want to reconstruct. For DAGER, it could become an issue when the embeddings of tokens in the reconstructed batch are linearly dependent. However, this is exceedingly rare for $b < d$ in practice, leading to only few false positives in the first filtering step, which we illustrate in Figure 2. Even if they do occur, they are handled correctly by our algorithm via a search procedure that is refined via span checks in the 2nd layer. With this, we are empirically able to filter all false positives, as any linear dependencies will have disappeared after the first transformer block.

Q4. On page 6, it is mentioned that the second layer filtering allows it to recover the exact 32 "starting tokens". What does it mean here "starting tokens"? Why is this an important issue to mention here?

We are happy to clarify this point! By “starting tokens” we refer to the first tokens of each of the to-be-reconstructed sequences, not any special tokens, such as [PAD] or [BOS]. The purpose of the experiment shown in Figure 2 is to demonstrate that the thresholds $\tau_1, \tau_2$ can take a wide range of values that will produce essentially the same correct solution. We mention the starting tokens in this context to simplify the presentation. The embeddings related to the 1st position are not related to any other tokens, and hence the experiment can isolate the two filtering steps without relying on any other knowledge. This is used to set the stage for the inductive progression to reconstruct the whole sequence for decoder-only models. We will add this clarification to the paper.

Q5. Would Filtering at higher layers (L>2) also help?

Yes, filtering at more layers would help reduce false positives. However, we observed that filtering at 2 layers is sufficient in non-noisy settings for exact reconstruction. There, further layers do not provide meaningful benefits. Therefore, we only considered two filtering layers to keep DAGER memory- and time-efficient. However, when the gradients are noisy as discussed in the response to Reviewer $\RL$’s Q6, additional layers provide significant benefits as they can be leveraged to average noise across different layers.

Conclusion

We hope to have been able to address all of the reviewer’s questions, are happy to answer any follow-up they might have, and look forward to their response.

References:

[A] Zhang, Zhuo, et al. "Fedlegal: The first real-world federated learning benchmark for legal nlp." Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2023.

$\newcommand{\RU}{\textcolor{blue}{Ujeg}}$We would like to thank reviewer $\RU$ for their very positive review, the provided insights and helpful recommendations. We are happy they found our experiments and proofs to be solid and our method innovative. Further, we are glad that the reviewer credits the ablation studies to be complete and is particularly happy that DAGER is capable of handling long sequences and large batch sizes. Next, we address all points raised by the reviewer.

Q1. Notation is heavy and hard to follow. Could authors add a notation table, or input before the algorithm so it is easier to follow?

Yes, we provide a notation table in the rebuttal PDF and will include it in the paper.

Q2. Can you show that DAGER works on more recent models, such as LLaMa-3?

Yes, we added experiments for both LLaMa-3.1 (8B) and LLaMa-3 (70B), demonstrating that DAGER achieves perfect reconstruction for both (see Table 5 in the attached PDF). A more elaborate description of our experiments and findings can be seen in the answer to Q1 in the main response.

Q3. Could the author discuss the usage of quantized models and the influence of LoRA for fine-tuning on reconstruction results?

Yes, we show that DAGER is directly applicable to 16-bit quantization, with 8-bit support for the setting we describe being part of active research. Further, we show that we work in the LoRA setting where the rank of the LoRA matrices is lower-bounded by $r>b$. A thorough description of our arguments and experiments can be found in Q3 in the main response.

Q4. Could the authors add some reconstructed text for comparison?

Yes, we show examples of reconstructed text for the Rotten Tomatoes dataset on a batch size of $B=1$ in the rebuttal. We choose $B=1$ as this is a practical limitation for the baseline. As can be observed in Table 3 in the rebuttal PDF, we perfectly reconstruct the entire sentence, while the best-performing baseline manages to only get a fraction of the sentence correct.

Conclusion

We hope to have been able to address all of the reviewer’s questions, are happy to answer any follow-up they might have, and look forward to their response.

$\newcommand{\RL}{\textcolor{red}{L4TG}}$We would like to thank reviewer $\RL$ for the very positive review and feedback. We are happy to read that the reviewer finds our work very smart, practical and novel. Further, we are glad they assessed our paper as easy to follow and our approach as highly effective. We thank the reviewer for giving us pointers on how to broaden the scale of DAGER. We now address all points raised by the reviewer.

Q1. What is the computational complexity of DAGER and does it scale to very large vocabulary sizes and long sequences?

We report both the runtimes (Table 5) and complexity (Theorem B.4) of DAGER in the Appendix of our paper. We highlight that the computational bottleneck for all practical parameters stems from the term $P^3 B^3 d^2$ or $P^4 B^3 d^2$ (in the RoPE case). As the vocabulary size is a linear component, it does not dominate the overall complexity of DAGER. We demonstrate the effectiveness on larger vocabularies by evaluating state-of-the-art models, such as LLaMa-3 70B and LLaMa-3.1 8B, which feature a vocabulary of 128,256 tokens, which is among the largest ones in existence. There, we observе only marginally increased runtimes.

Please refer to Q1 in the main response and Table 5 in the rebuttal PDF for a description of the results and requirements.

Q2. How does DAGER relate to current state-of-the-art methods in gradient leakage in the text domain?

DAGER assumes the honest-but-curious setting, where the server is non-malicious, unlike the Decepticons [28] and Panning for Gold [29] attack, meaning we do not need to change the model weights sent to clients. Most prior honest-but-curious methods, such as DLG [6], TAG [10], LAMP[11] and Jianwei et al. [31] rely on relaxing the discrete optimization problem to a continuous one that has no guarantees of convergence. One exception is APRIL [23], which analytically recovers the inputs but doesn’t scale to B>1. DAGER, however, can handle B=128. Another one is FILM [22], which requires access to the gradients of a batch at any iteration and, similarly to LAMP, requires a strong language prior, unlike DAGER.

Further, we do not require knowledge of the labels for different sentences in the batch in the sentiment analysis case, in contrast to TAG and LAMP, and we work with almost any reasonable loss function (see Q2 in the main response) unlike Flat-Chat [32]. Finally, in contrast to APRIL, FILM and Decepticons, we do not require the gradient of the embedding layer, making our setting significantly harder. We will clarify these points in the revised version of the paper.

Q3. Would DAGER work in a decentralized FL setting with a client acting as the attacker?

While it is difficult to generalise under the vastly different decentralized FL (DFL) protocols, we attempt to provide a unified answer. The crucial part under most DFL protocols, such as IPLS [A], ProxyFL [B], or TrustedDFL [C], is that not gradients but model weights are shared between clients. We see 2 difficulties here - having multiple training steps before a model is shared, and having model weights that are aggregated across clients whose updates the attacker may not receive.

The former is tightly associated with the FedAvg setting, which we have successfully applied to DAGER (see Table 4 in the rebuttal PDF). The latter is, however, a protocol-specific task, and is an interesting direction for future work.

Q4. In the Centralized FL environment, would asynchronous training affect the performance of DAGER or render it infeasible?

In centralised asynchronous FL the server and clients exchange updates asynchronously with clients sometimes creating updates based on more than one model shared by the server. Crucially, an honest-but-curious server can keep track of all updates sent to a client. This means that while a client’s gradients can be computed across several models, the models are all known to the server (as in the protocol described in [D]). Further, assuming linear aggregation, the rank of the gradient is still upper-bounded by the number of tokens used to compute it. The same formula from Theorem 3.1 still applies with $\mathbf{X}$’s coming from different models. To this end, the server will need to apply DAGER for tokens propagated through all possible client models.

Q5. There's a small error in Figure 1 - AA_ should not be in the span and BA_ should be.

Yes, thank you for spotting it! We will fix it.

Q6. How does DAGER perform against noise (DP) or malicious users? Can it be amended to combat this?

It has been shown [E] that DP provides provable guarantees for protecting privacy. Nevertheless, we provide promising initial results, where we use deeper layers to mitigate the effect of noise. The crucial assumption here is that noise is independent across layers. We apply DAGER on different noise levels with $B=1$ on the Rotten Tomatoes dataset utilising the GPT-2 model. The results can be found in Table 2 in the rebuttal PDF. We believe that there are numerous improvements one could make, but leave these for future work.

It is unclear to us what the reviewer refers to as “malicious user” and how input reconstruction could help there. We kindly ask for some clarification and are more than happy to address the question.

Conclusion

We hope to have been able to address all of the reviewer’s questions, are happy to answer any follow-up they might have, and look forward to their response.

References:

[A] Pappas, C., et al. "Ipls: A framework for decentralized federated learning." 2021 IFIP.

[B] Kalra, S., et al. "Proxyfl: decentralized federated learning through proxy model sharing." (2021).

[C] Gholami, A., et al. "Trusted decentralized federated learning." CCNC, 2022.

[D] Chen, Y., et al. "Asynchronous online federated learning for edge devices with non-iid data." Big Data. 2020.

[E] Abadi, M., et al. "Deep learning with differential privacy." ACM SIGSAC. 2016.