DAGER: Exact Gradient Inversion for Large Language Models

Federated learning works by aggregating locally computed gradients from multiple clients, thus enabling collaborative training without sharing private client data. However, prior work has shown that the data can actually be recovered by the server using so-called gradient inversion attacks. While these attacks perform well when applied on images, they are limited in the text domain and only permit approximate reconstruction of small batches and short input sequences. In this work, we propose DAGER, the first algorithm to recover whole batches of input text exactly. DAGER leverages the low-rank structure of self-attention layer gradients and the discrete nature of token embeddings to efficiently check if a given token sequence is part of the client data. We use this check to exactly recover full batches in the honest-but-curious setting without any prior on the data for both encoder- and decoder-based architectures using exhaustive heuristic search and a greedy approach, respectively. We provide an efficient GPU implementation of DAGER and show experimentally that it recovers full batches of size up to 128 on large language models (LLMs), beating prior attacks in speed (20x at same batch size), scalability (10x larger batches), and reconstruction quality (ROUGE-1/2 > 0.99).

Paper

Similar papers

Peer review

Reviewer ju5t6/10 · confidence 4/52024-07-03

Summary

The authors proposed a method to recover user training data from gradients in federated learning. The key observation is if a linear network component exists in the overall neural network, the gradient of the linear component parameter can be shown as a linear combination of the input to this linear component. Taking the the first layer of the transformer, the input token embeddings can then be recovered, assuming the number of tokens involved is smaller than the dimension of the embedding space. Extending the idea to the second layer helps to recover the whole training sequence. The authors show empirically the proposed method can significantly improve the state of the art.

Strengths

The idea is rather straightforward and effective. The numerical improvement over the state of the art is impressive.

Weaknesses

1. The theory development is weak. The idea is rather straightforward, however, the authors need to lay down more carefully the technical assumptions. For example, is there any requirement on the loss function, i.e., MSE or cross-entropy? Are there any requirements on the differentiability? For ReLU, the loss function is sub-differentiable, and how does this impact the method? 2. The requirement that b<d needs better justification. The embedding space is usually not very large (e.g, 512, 768, or 1024). However, the number of tokens used by a single client in a single training step (or before synchronizing with the central server) can be considerably larger than this. The authors argue the condition is satisfied for reasonable input lengths and batch sizes, which I found to be less than convincing. 3. Conceptually it is not clear why the method in Algorithm 1 works so well, and some better explanation is needed. The embedded tokens are not linearly independent, and there can be significant false positives in this procedure, yet the authors observed otherwise. A more careful analysis is warranted.

Questions

Some questions are as given above. Additionally, 1. 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? 2. Would higher layers (L>2) also help? Why is this not considered?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes, it was discussed in the conclusion, particularly when B is large.

Reviewer Ujeg7/10 · confidence 3/52024-07-13

Summary

This is the first paper to use low rank decomposition to attack the gradient of the self-attention layer to extract information for LLM. It also provides a fast algorithm to recover the correct token first, then the sequence.

Strengths

Solid experiments and math proof. It’s a good innovative finding, especially using low rank decomposition on self-attention part. Especially, it focuses on long sequences like 4096 and large batches for 128 on decode-base structure. The authors also provide complete ablation studies and detailed experiments in many aspects, like model size and model parameters. Its much better than a lot of papers in the language field.

Weaknesses

Notation is heavy and hard to follow. Could authors add a notation table, or input before algorithm so it is easier to follow? Some people might argue it not use latest model but working on llama3 could be resource consuming. It might be a weaknesses of the author's work, but it could also be a future direction. If possible, could the author discuss a little bit about the usage of quantized models and the influence of lora for fine-tuning regarding reconstruction results? Also, could authors add some reconstructed text for comparison?

Questions

see weakness above

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

see weakness above

Reviewer L4TG7/10 · confidence 2/52024-07-13

Summary

In this paper, the authors propose the DAGER algorithm which leverages the low-rankness of self-attention layer gradients in order to filter out incorrect embedding vectors and recursively reconstruct true input text. DAGER works within the Centralized Federated Learning setting, where the server is honest in aggregating updates and performing model updates, but is curious to try to recover personal information from each client.

Strengths

1. A very smart/practical and, to the best of my very limited knowledge in this area, novel approach to reconstruct text by disregarding all embeddings which do not align with the low rank gradients received by the server. 2. The empirical results demonstrate that this method is highly effective, and greatly outperforms all other baselines. 3. Paper is well-written, has good diagrams and figures, and is easy to follow for someone not well-versed in this area (attacks & LLMs).

Weaknesses

1. I'm curious about the computational complexity of DAGER. It seems that the search space can be extremely large (even when heuristics are used), and recovering individual tokens could take awhile. 2. (minor) In the related works it would be nice to go deeper into research within the gradient leakage in the text domain. A specific breakdown of how this paper differs from the current methods.

Questions

Would this be able to work within the decentralized FL environment (i.e. could a single client, acting as the aggregator, decipher other agent text)? In the Centralized FL environment, would asynchronous training affect the performance of DAGER or render it infeasible (asynchronous, stale updates)? There's a small error in Figure 1 in Sequence recovery (middle part, 2nd sequence), AA_ should not be in the span and BA_ should be? In the diagram the opposite is true. How does DAGER perform against noise (DP) or malicious users? Can it be ammended to combat this? What happens if the vocabulary size is extremely large? Won't this process take a very long time (i.e. does it scale with |V|)? How expensive is this enumeration process? Especially if the sequence is long? Very cool insight about size of model and how that allows gradients to be reconstructed easier because they're low rank since d is increasing so rapidly. Bigger is not always better when it comes to privacy!

Rating

7

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes.

Reviewer ju5t2024-08-12

Thanks for the response. I feel the authors will be able to make the expected improvement/clarification based on these questions/responses. I'll keep the score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC