Differentially Private In-Context Learning with Nearest Neighbor Search

AI summary

Overview

• Research area: Privacy-preserving machine learning — specifically differentially private in-context learning (DP-ICL) for large language models, combined with nearest-neighbor (kNN) similarity search and individual differential privacy accounting.
• Technical level: Advanced. The paper relies on Rényi differential privacy (RDP), δ-approximate RDP, privacy filters, adaptive composition, propose-test-release (PTR), and report-noisy-max mechanisms.
• Scope: The paper proposes a method (DP-KSA-kNN) that replaces random subsampling of in-context demonstrations with nearest-neighbor retrieval, using individual δ-approximate RDP privacy filters for accounting, and evaluates it on private text classification and private document question answering.

What This Paper Is About

Existing differentially private in-context learning methods privatize how a language model aggregates or votes over demonstrations, but they retrieve those demonstrations by random subsampling. The paper argues this ignores a standard component of modern LLM pipelines — similarity search — and that randomly sampled examples can be unrelated to the query, increasing prediction uncertainty and potentially performing worse than zero-shot prediction. The goal is to plug nearest-neighbor retrieval of the most query-relevant examples into an existing DP-ICL framework, while still respecting a central differential privacy budget through per-example privacy filters.

Key Contributions

1. kNN retrieval inside DP-ICL. The authors integrate nearest-neighbor search into the DP-ICL framework of Wu et al. (2024). Prompts are composed with the k-nearest neighbors of each test point instead of randomly sampled examples, as in the baselines by Wu et al. (2024) and Tang et al. (2024). This is done via privacy filters (Feldman and Zrnic, 2021).
2. A theoretical result on adaptive RDP filters. The paper provides a fully adaptive δ-approximate RDP analysis of individual RDP filters (Theorem 1), described as a generalization of the RDP filtering result in Feldman and Zrnic (2021, Thm. 4.5) and the δ-approximate zCDP filtering result in Whitehouse et al. (2022, Thm. 1).
3. A stability requirement for retrieval. The retrieval function is required to be stable under single-element change: neighbors in the dataset must produce retrieved sets differing by at most two elements, which bounds the sensitivity of the aggregation. The authors use the FLAT index because exhaustive search trivially satisfies this property.
4. Empirical evaluation on two task families with two LLM families. Experiments cover text classification (AGNews, TREC) and document question answering (Federated DocVQA, SQuAD v1.1) using OPT-1.3B, Llama3.3-70B-it, and Gemini-1.5-flash-8B.

Main Findings

• Substantial gains on classification benchmarks: The abstract and introduction state the method outperforms existing baselines by a substantial margin across all evaluated benchmarks and achieves more favorable privacy-utility trade-offs. The results are plotted in Figure 1 rather than reported as a numeric table.
• Experimental setup matters for interpretation: Text classification used 200 randomly sampled test samples per dataset, δ = 10⁻⁵, results averaged over 5 independent runs with error bars at 1.96 times the standard deviation (asymptotic 95% confidence interval).
• Low-ε degradation on TREC: On TREC, performance deteriorates as ε decreases toward 0.5. The authors attribute this to the Gaussian mechanism's sensitivity of √2: at ε = 0.5 the required noise scale is approximately 10, which equals the number of shards (number of votes), significantly randomizing predictions. They note using more shards could remedy this.
• Zero-shot caveat for OPT-1.3B: Zero-shot results were excluded from Figure 1 because the zero-shot baseline reaches approximately only 58% test accuracy for AGNews, while TREC results are close to random guessing. The authors note baseline and zero-shot accuracies are similar to Wu et al. (2024), which uses GPT-3 Babbage with approximately 1.3B parameters.
• Document QA: Figure 2 shows 4-shot ICL with shard sizes 10 and 20 across several ε values, with the same 100 randomly sampled test queries used for all methods and ε values. Points at ε = ∞ correspond to the non-private versions of KSA and KSA-kNN. The authors emphasize that outperforming the 0-shot baseline (computed with the same number of shards) is important to justify using private demonstrations.
• Model comparison: The main high-level observation across both figures is that most metrics have higher values for the Llama model compared to Gemini.
• Differing ε sensitivity: DP-KSA remains less sensitive to ε, whereas DP-KSA-kNN improves in many cases, especially for high ε. The authors state DP-KSA-kNN is superior to the DP-KSA baseline.
• No subsampling amplification: The approach does not benefit from privacy amplification by subsampling, yet still achieves better privacy-utility trade-offs by tolerating higher noise in the aggregation step.
• Baseline attribution note: The conclusion states the method outperforms the zero-shot and the DP baseline method by Tang et al. (2024); the experimental comparisons in the main text are against DP-KSA from Wu et al. (2024).

Methodology in Plain English

The pipeline works like this. There is a private database of example records (for classification, an article plus label; for QA, a paragraph or OCR'd invoice text plus question and answer). For each incoming query, the system first does a nearest-neighbor search in embedding space to pull out the k most similar examples, rather than picking examples at random.

Those retrieved examples are split into M disjoint batches, and each batch is used to build a separate in-context prompt. The language model answers each prompt independently, producing a set of responses. For text classification, the responses are class labels, so the system forms a histogram of label votes and uses report-noisy-max with Gaussian noise to pick the winner. For document question answering, the answers are free text, so the system instead builds a histogram of token frequencies across all responses and privately releases the top-K keywords using propose-test-release; those keywords are then fed into a follow-up prompt that asks the model for a concise final answer.

The privacy bookkeeping is the delicate part. Because the retrieval step is data-dependent, the authors use an individual privacy filter that tracks each data element's cumulative privacy loss across queries. At each step, the filter computes per-example RDP costs for the current mechanism, updates the set of "active" data elements whose cumulative ε and δ remain within budget (ε_max, δ_max), and drops any element about to exceed it. To keep the sensitivity bounded, the retrieval function must be stable: changing one database element may change the retrieved set by at most two elements, which the exhaustive FLAT index guarantees. Embeddings are produced with the "all-MiniLM-L6-v2" model, giving unit-length 384-dimensional vectors, and FAISS is used to build the FLAT indices.

Why This Matters

Impact on research. The paper connects two previously separate lines: private prediction methods that use kNN with individual privacy accounting (e.g., Zhu et al., 2023), and DP-ICL methods built on private aggregation of LLM outputs. It shows that similarity search — a standard, easily pluggable component of retrieval-augmented generation — can be inserted into DP-ICL without abandoning rigorous privacy guarantees, and it supplies a more general adaptive RDP filtering theorem to make that accounting sound.

Real-world applications:

• Private question answering over confidential business documents such as invoices containing payer/payee names, amounts, and purpose (the Federated DocVQA setting used in the paper).
• Classification of sensitive user text, such as news categorization (AGNews) or question-type classification (TREC), without exposing individual training examples.
• Retrieval-augmented generation systems built on private corpora (legal, medical, or enterprise knowledge bases) where demonstrations must not leak their contents.
• API-only LLM deployments where no logit access is available, since the method only needs prompt access to the model.

Industry relevance. The authors are affiliated with Nokia Bell Labs and Nokia, and the method is designed to plug into existing ICL/RAG pipelines with standard libraries (FAISS) and off-the-shelf models (OPT-1.3B, Llama3.3-70B-it, Gemini-1.5-flash-8B). Because it requires only prompt-level access and no model weight updates, it is practical for commercial LLM services handling confidential data.

Future Directions

• Approximate indexing. The current method uses only the exhaustive FLAT index for simplicity. Extending to IVF or HNSW is described as a compelling avenue, and the authors suggest DP k-means methods (Chang et al., 2021) could implement IVF search while still meeting the stability requirement, at an additional privacy cost.
• Removing the low-ε weakness. The TREC degradation at low ε could be addressed by increasing the number of shards, which the authors identify as a remedy.
• Adaptive failure probabilities. The paper fixes the iteration-wise failure probability δᵢ when using DP-KSA, and notes it could instead be chosen adaptively.
• Broader indexing and retrieval design. The conclusion calls for DP-ICL solutions built on alternative sample indexing and retrieval methods, such as hierarchical clustering like k-means or hierarchical navigable small worlds (HNSW).

Target Audience

This paper is best suited to machine learning researchers and engineers working on privacy-preserving LLM systems, particularly those familiar with differential privacy and Rényi differential privacy. It also benefits practitioners building retrieval-augmented generation or in-context learning pipelines over sensitive data, and privacy researchers interested in individual privacy accounting and privacy filters for adaptive data-dependent mechanisms. The theory in Section 2.3 and the appendix assumes comfort with DP definitions, so readers seeking only applied guidance will want to focus on the introduction, experiments, and conclusions.

Authors’ abstract

Differentially private in-context learning (DP-ICL) has recently become an active research topic due to the inherent privacy risks of in-context learning. However, existing approaches overlook a critical component of modern large language model (LLM) pipelines: the similarity search used to retrieve relevant context data. In this work, we introduce a DP framework for in-context learning that integrates nearest neighbor search of relevant examples in a privacy-aware manner. Our method outperforms existing baselines by a substantial margin across all evaluated benchmarks, achieving more favorable privacy-utility trade-offs. To achieve this, we employ nearest neighbor retrieval from a database of context data, combined with a privacy filter that tracks the cumulative privacy cost of selected samples to ensure adherence to a central differential privacy budget. Experimental results on text classification and document question answering show a clear advantage of the proposed method over existing baselines.