DAGER: Exact Gradient Inversion for Large Language Models

Overview

• Federated learning allows for collaborative training without sharing private client data, but prior work has shown that the data can be recovered using gradient inversion attacks.
• While these attacks work well on images, they have limitations in the text domain and can only approximately reconstruct small batches and short input sequences.
• This paper proposes DAGER, the first algorithm to exactly recover whole batches of input text.

Plain English Explanation

Federated learning is a way for multiple computers or devices to work together on a machine learning task without having to share their private data. Instead of sending the actual data, the computers send only the changes, or "gradients," that they want to make to the machine learning model. The server then combines these gradients to update the model.

However, prior research has shown that the server can actually use these gradients to figure out what the original data was. This is called a "gradient inversion attack." While these attacks work well on things like images, they have limitations when it comes to text data.

In this paper, the researchers propose a new algorithm called DAGER that can exactly recover whole batches of input text, even for large language models. DAGER takes advantage of the way self-attention layers and word embeddings work to efficiently search for the original text that matches the gradients.

Technical Explanation

The key innovation in DAGER is leveraging 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. This allows DAGER to exactly recover full batches of text, rather than just approximate reconstructions of small fragments.

The paper describes two versions of DAGER - one that uses an exhaustive heuristic search for encoder-based architectures, and one that uses a greedy approach for decoder-based architectures like large language models (LLMs). The researchers provide an efficient GPU implementation of DAGER and show that it outperforms prior attacks in terms of speed (20x faster at the same batch size), scalability (able to handle 10x larger batches), and reconstruction quality (ROUGE-1/2 > 0.99).

The key technical insights come from exploiting the low-rank structure of self-attention gradients and the discrete nature of token embeddings to efficiently search the space of possible text sequences.

Critical Analysis

While DAGER represents a significant advance in gradient inversion attacks on text data, the paper acknowledges some important limitations and areas for further research. For example, the attacks are demonstrated in an "honest-but-curious" setting, where the server follows the protocol but tries to recover the data. More adversarial threat models, where the server actively tries to subvert the protocol, are an important area for future work.

Additionally, the paper only considers relatively short input sequences and small batch sizes, compared to the scale of real-world language models and datasets. Dealing with this uncertainty and scaling up the attacks will be crucial for understanding the real-world implications of gradient leakage in federated learning.

Overall, this paper makes an important contribution by demonstrating that even text data, which was believed to be more resistant to gradient inversion, can be vulnerable to attacks under certain conditions. This highlights the need for continued research into secure federated learning protocols that can protect against these types of threats.

Conclusion

This paper introduces DAGER, the first algorithm capable of exactly recovering whole batches of input text from the gradients shared in a federated learning setting. By exploiting the low-rank structure of self-attention gradients and the discrete nature of token embeddings, DAGER outperforms prior attacks in terms of speed, scalability, and reconstruction quality.

While DAGER represents a significant advance, the paper also acknowledges important limitations and areas for future research, such as exploring more adversarial threat models and scaling up the attacks to handle larger language models and datasets. Overall, this work highlights the ongoing challenge of protecting the privacy of user data in federated learning systems.