Don’t forget private retrieval: distributed private similarity search for large language models (2024)

Abstract: While the flexible capabilities of large language models (LLMs) allow them to answer a range of queries based on existing learned knowledge, information retrieval to augment generation is an important tool to allow LLMs to answer questions on information not included in pre-training data. Such private information is increasingly being generated in a wide array of distributed contexts by organizations and individuals. Performing such information retrieval using neural embeddings of queries and documents always leaked information about queries and database content unless both were stored locally. We present Private Retrieval Augmented Generation (PRAG), an approach that uses multi-party computation (MPC) to securely transmit queries to a distributed set of servers containing a privately constructed database to return top-k and approximate top-k documents. This is a first-of-its-kind approach to dense information retrieval that ensures no server observes a client’s query or can see the database content. The approach introduces a novel MPC friendly protocol for inverted file approximate search (IVF) that allows for fast document search over distributed and private data in sublinear communication complexity. This work presents new avenues through which data for use in LLMs can be accessed and used without needing to centralize or forgo privacy.

## introduction heavily pre-trained and fine-tuned large language models (llms) have demonstrated exceptional performance on zero-shot @xcite and few-shot tasks @xcite . the ability of these models to generalize, combined with their costly pretraining, has shifted the focus from training ad-hoc models to perform specific tasks to utilizing these general-purpose foundational models for a wide variety of use-cases @xcite . these pre-trained models lack knowledge of private contexts or recent events. to provide these llms with up-to-date or relevant information, methods such as retrieval augmented generation (rag) @xcite @xcite are used to include external information into a generation process without needing fine-tuning on new data. this process allows llms to first query an external data source, retrieve relevant information (with respect to a given prompt), and then use both the prompt and the retrieved data as input to the inference phase of the llm. similar to the problem of federated learning @xcite , it is valuable to aggregate sensitive data from multiple (perhaps many) data owners. to do that, each party should be able to guarantee that their own private data remains private even when it is utilized. on the other hand, model users should be able to query these data from many data owners without needing to share what questions they are asking. in this work we argue that llms require a new model for sharing data for ai tasks. compared to federated learning, which focuses on the training phase, llms should focus on the (i) retrieval phase; (ii) inference phase. guaranteeing privacy of both the query and any private documents residing in the retrieval database require that both phases utilize privacy-preserving techniques and are chained together. alas, to the best of our knowledge all existing works only tackle the llm inference problem @xcite @xcite , but provide no secure solution when retrieval is involved. in this work, we close this gap by introducing private retrieval augmented generation (prag). prag allows users to privately search a database, which in itself is private, then send the augmented query privately to any secure (or otherwise trusted) llm, creating an end-to-end secure solution. our approach and contributions. in this paper, we propose private retrieval augmented generation (prag), a secure approach to augment neural information retrieval that hides both query vectors and the retrieval database. we use a retrieval database split across a set of servers, and we ensure data remains private by using secure multi-party computation (mpc) techniques. to the best of our knowledge, we are the first to consider the problem of secure distributed retrieval in the context of llms, and more broadly, are the first to propose a solution for private similarity search that can protect both the query and a secret-shared (or encrypted) database. this approach can be deployed with any standard neural information retrieval (ir) embedding model to augment distance calculations (e.g., cosine, dot, euclidean) and topk retrieval over federated vector stores, scaling to medium-size databases with very little accuracy loss (99% accuracy on real data). we further scale the approach to much larger databases using an approximate k-nearestneighbors approach inside mpc, replicating the accuracy of the state of the art in approximate retrieval using a first-of-its kind inverted files index inside mpc, providing significant speed improvements for retrieval. our approach provides both theoretical and empirical improvements of value. we achieve constant communication on the client's side and sublinear communication on the servers' side --the bottleneck in mpc approaches. this work is the first ir approach to work across more than two servers with minimal additional costs. we further present a 'leaky' version of the protocol that allows for partial privacy of queries under a privacy budget with significant improvements to speed. we evaluate prag across a range of data distributions, both real and synthetic, to show it broadly maintains the performance characteristics of nonsecure ir approaches. we provide a pytorch-native implementation of our system using the crypten mpc engine foot_0 . ## methods in this section, we present the private retrieval augment generation (prag) framework. the method builds from secret sharing and mpc friendly exact top-k calculations to a new mpc design of an inverted file index for efficient approximate top-k calculation. a visual high-level overview of this design and its usage with a client llm querier is shown in figure 1 . ## experiments to demonstrate the performance of these models we run a series of experiments on both synthetic and real data to determine performance properties of the implementations of these methods above. we benchmark the retrieval accuracy and speed across a range of embedding sizes (256 to 8192), synthetic embedding distributions (n (0, 0.05), n (0, 1), u (-1, 1), binary), distance functions (cosine, dot product, euclidean), top-k values, ivf parameters, and database sizes. we perform mpc experiments on a single 2.2ghz intel xeon silver cpu using crypten's built-in communication code to spawn processes for each server. further to this, we test the approaches on retrieval of real neural embedding datasets from beir @xcite using the same environment, this collection of datasets uses a range of textual document types and sizes, all of which we use a standard off-the-shelf embedding on. while there are several parallelization improvements that can be made locally within each server for mpc, our implementations of each algorithm above remain unoptimized. ## related work drawing on the ideas in private federated learning, we can maintain privacy when doing public queries @xcite and move beyond in-context learning @xcite . we bring privacy to this idea through augmenting existing non-private retrieval methods, ranging from exact search on small datasets to large scale approximate retrieval @xcite . while several other works have examined the problem of secure similarity search @xcite @xcite @xcite @xcite , to the best of our knowl-edge we are the first to examine a model where the database is secret shared as well, and where an arbitrary number of servers and database owners can be supported. a comparison to the state-of-theart protocols @xcite is available in these approaches can augment other pieces of privacy-first ml infrastructure from fully secure llm inference @xcite and federated or privacy preserving k-means clustering @xcite . we choose to focus on mpc techniques in this paper, as opposed to secure retrieval schemes that rely trusted execution environments (tees) @xcite @xcite , as tees have been known to suffer from privacybreaching attacks. ## conclusion we introduced prag, a novel approach for secure, distributed information retrieval for large language models. prag uniquely safeguards both query vectors and a multi-owner database using multiparty computation (mpc). key contributions include an mpc-friendly protocol for inverted file approximate search, allowing for rapid document retrieval with sublinear communication complexity; analysis of exact search performance on language embeddings; and a version of the protocol that offers a trade-off between speed and partial privacy, under a predefined privacy budget. these tools allow for a new mechanism of neural information retrieval, which when combined with secure inference of llms, is a stepping stone towards fully secure foundation model agent pipelines. however, much like secure execution of llms, the approach put forward here has significant computational costs and speed limitations, especially for large databases and high accuracy demands. future work should explore optimizing communication costs, expanding beyond a semi-honest adversary, and integrating prag into larger secure machine learning frameworks.