Selective Pre-training for Private Fine-tuning

TL;DR

This work tackles the challenge of training small, domain-specific language models under differential privacy with strict inference-time constraints. It introduces selective pre-training, a data-selection approach that uses a DP domain classifier to curate a public pre-training subset aligned with the private target distribution, enabling effective private fine-tuning on constrained models. Empirical results on Enron and GLUE show that selective pre-training yields state-of-the-art DP performance for small models and can let smaller models match larger non-private baselines, with private learning benefiting more from data quality than non-private setups. The method promises practical privacy-preserving model compression and reduced inference costs, while noting DP guarantees primarily apply to private data and that pre-training from scratch incurs additional compute cost.

Abstract

Text prediction models, when used in applications like email clients or word processors, must protect user data privacy and adhere to model size constraints. These constraints are crucial to meet memory and inference time requirements, as well as to reduce inference costs. Building small, fast, and private domain-specific language models is a thriving area of research. In this work, we show that a careful pre-training on a \emph{subset} of the public dataset that is guided by the private dataset is crucial to train small language models with differential privacy. On standard benchmarks, small models trained with our new framework achieve state-of-the-art performance. In addition to performance improvements, our results demonstrate that smaller models, through careful pre-training and private fine-tuning, can match the performance of much larger models that do not have access to private data. This underscores the potential of private learning for model compression and enhanced efficiency.

Figures (10)

• Figure 1: The proposed framework for training a small and domain-specific model with differential privacy (DP). More details on the process of training the domain classifier and the selection of public data can be found in Figure \ref{['fig:illustration_selection']}. We use the method in AbadiCGMMTZ16 for training models with DP.
• Figure 2: A representative result from our findings. We plot perplexity and top-1 next word accuracy of GPT models on the test set of the Enron email dataset enron. Overall privacy budget is $(\varepsilon=7.3, \delta=1\times 10^{-7})$. The dashed line shows the zero-shot performance of GPT2-XL with 1.5 billion parameters. The figure shows that our framework yields clear improvements in both perplexity and next-token prediction accuracy, which can significantly improve the overall model behavior.
• Figure 3: The process of training the domain classifier and the selection of large-scale public data.
• Figure 4: The 100 most frequent nouns in Enron email, OpenWebText, or a selected subset of OpenWebText (10%). A larger font size indicates that the word appears more frequently. Green words are the 100 most frequent nouns in Enron Email. OpenWebText and selected OpenWebText have 28 and 39 words, respectively, that are among the 100 most frequent nouns in Enron Email.
• Figure 5: Perplexity and top-1 next word accuracy of GPT models on the test set of the Enron email dataset. The models are trained without DP. Selective pre-training still improves over standard pre-training, however, the improvements are smaller compared to private learning.

Theorems & Definitions (1)

• Definition 1: $(\epsilon,\delta)$-Differential Privacy (DP) dwork2006calibrating