Training Text-to-Text Transformers with Privacy Guarantees (2022)

Abstract: Recent advances in NLP often stem from large transformer-based pre-trained models, which rapidly grow in size and use more and more training data. Such models are often released to the public so that end users can fine-tune them on a task dataset. While it is common to treat pre-training data as public, it may still contain personally identifiable information (PII), such as names, phone numbers, and copyrighted material. Recent findings show that the capacity of these models allows them to memorize parts of the training data, and suggest differentially private (DP) training as a potential mitigation. While there is recent work on DP fine-tuning of NLP models, the effects of DP pre-training are less well understood: it is not clear how downstream performance is affected by DP pre-training, and whether DP pre-training mitigates some of the memorization concerns. We focus on T5 and show that by using recent advances in JAX and XLA we can train models with DP that do not suffer a large drop in pre-training utility, nor in training speed, and can still be fine-tuned to high accuracies on downstream tasks (e.g. GLUE). Moreover, we show that T5’s span corruption is a good defense against data memorization.

## introduction recent advances in natural language processing tasks are largely due to introduction of large transformer-based models trained on large amounts of data. models such as gpt-2 @xcite and t5 @xcite have billions of parameters and are trained on hundreds of gigabytes of mostly uncurated public crawl data. these models are often released as modifiable checkpoints, and the end users have the ability to fine-tune these models to their final tasks using an often more limited amount of data and compute. while pre-training datasets are typically treated as public, their sheer size makes them difficult to curate or scrutinize @xcite . moreover, such public datasets (e.g., web crawls) likely contain private information @xcite , e.g., data erroneously released to the web or copyrighted text. the capacity of recent models makes it possible for them to memorize parts of the training data @xcite , even after subsequent fine-tuning, and poses risks to the owners of pre-trained language models. in this work, we focus on a potential mitigation: making the model fully private using differential privacy (dp). we focus on t5 @xcite and explore how well dp mitigates privacy risks and how it affects pre-training and downstream performance. our contributions are as follows: in this section, we discuss what privacy in ml means ( §2.1), followed by overview of private training ( §2.2) and privacy in language models ( §2.3). ## implementing a fully private t5 we focus on t5 @xcite , a popular encoder-decoder. it uses a slightly modified transformer architecture @xcite and both the input and output is a sequence of tokens, as tokenized by sentencepiece @xcite . t5 is a good model to focus on, since it can be used for many input-output tasks, is trained on a large public crawl data set, is publicly available, and has been shown to have excellent performance on subsequent fine-tuning tasks. what we are protecting. we use the dp definition, so we provide protection at the level of a training instance. for encoder-decoders like t5, that means a pair of input and output sequences. importantly, if the same training example is repeated multiple times in the training data, the level of protection for such an example will be smaller. ## experiments hyperparameters. our experiments in the main text use t5 small, which has 6 encoder layers, 6 decoder layers, 8 64-dimensional heads, embedding dimension of 512, mlp dimension of 2048. we chose t5 small since it is relatively fast to train and produces results comparable with that of larger models (see appendix a for a discussion of the effect of the model size on pre-train and fine-tuning performance). we use adafactor @xcite with learning rate 0.5, decay rate 0.8, warm-up 1000, and rsqrt learning rate decay. datasets. we use the colossal clean crawled corpus (c4; @xcite as a pre-training task and look into the original "prefix" unsupervised training objective, that predicts next tokens given the context and the span corruption training objective @xcite , 2020, §3.3.4), where randomly removed spans of the input are predicted. we use 512 tokens as input/context and attempt to predict 114 target tokens. to evaluate fine-tuning performance, we utilize glue datasets @xcite that allows to evaluate model performance accross a range of nlu tasks. ablations. we look separately into the effect of dp-sentencepiece and dp-train on memorization and pretraining and subsequent fine-tuning perfor- pre-training and fine-tuning performance. it is known that dp-training hurts the utility (e.g., accuracy) of models. however, the common scenario in nlp is that models are pre-trained on some data and then subsequently fine-tuned for the end task. we look into private pre-training (contrary to the majority of the papers which look into private finetuning of a publicly pre-trained model), and we hope that (public) fine-tuning such privately pretrained models on public data provides the same utility as publicly pre-trained models. for these experiments, we train t5 models with the span corruption objective with a batch size of 8192 for 100k steps, and fine-tune on glue for 150k additional steps with a (standard) 128 batch size. the batch size of 8192 was chosen for pre-training since it provides good performance for dp-train. t5 without dp-training trains with approximately the same performance using a batch size of 128, however for dp-training it is known that the batch size should be increased significantly in order to get reasonable performance (see appendix d). for a fair comparison we use the same batch size for the baseline and dp-t5 variants. 2 another alternative is to tune hyper-parameters (for both baseline and dp variants) and compare the best possible models, however since tuning parameters for dp will change the guarantees, we don't go this route. we use an initial learning rate of 0.5 (both baseline and dp-t5 variants pre-training) and weight decay, and train with 64 cores. for dp-training, please refer to appendix 5 for details on noise, clipping norm and . for the full dp t5 model, we use a dp-sp unigram model trained on c4 @xmath0.17 (see additionally, appendix d discusses additional modifications that can be further explored to minimize the utility drop due to dp-train. testing for memorization. it is expected that dp (both dp-training and dp-sp) should reduce data memorization of the models. to verify that, we conduct an evaluation similar to @xcite and @xcite . in particular, we train models on c4 and attempt to "extract" the training data by providing an input prefix and allowing model to generate the rest of the sequence. we use 512 tokens prefix length, similar to the input length that was used for training. when the model generates the output, we calculate the exact match on a per-instance basis and report the average across all instances in the dataset. exact match is different from per-token accuracy: it is either 0 or 1 for each prefix, depending on whether the model generated the exact target sequence or not, whereas token accuracy would report how many predicted tokens match the target tokens for each instance. exact match allows us to gauge what fraction of instances was output verbatim by the model, serving as a useful metric for memorization. token-level accuracy serves as an additional metric; while getting some tokens right does not guarantee that memorization occurred, high values of token-level accuracy would indicate that some tokens generated by the model words might have matched those from the target exactly. finally, we also report median edit distance between predicted and target tokens, averaged across all instances in the data. this metric also serves as indirect way of measuring memorization. the difference between our setup and that of @xcite and @xcite is that our model is an encoder-decoder, as opposed to gpt-2 which is a decoder only. additionally, on top of using just next word prediction as a training objective (referred to as prefix training), we get to experiment with the span corruption objective. we test for memorization on the the same four c4-based datasets as @xcite : we would expect the most memorization to be exhibited on train dup, and the least memorization to be present on valid unique data. ## conclusion while the majority of recent work looks into private fine-tuning of pre-trained nlp models, we investigated how private pre-training of a model on a large corpus of data affects its pre-training and subsequent fine-tuning performance, as well as how much memorization such privately pre-trained models exhibit. we worked with t5, a transformerbased encoder-decoder, and demonstrated how to achieve a fully private t5 version by introducing dp-sentencepiece to train a differentially private subword tokenizer, and implementing dp-training for the actual pre-training. we leveraged recent advances in jax to do so without incurring a large performance hit in terms of training speed. our results show that both dp-sentencepiece and dp-training contribute to reducing memorization, but that the latter has the largest effect. moreover, we demonstrated that the span corruption task from @xcite also effectively mitigates memorization, which isn't the case for the next token prediction objective. we also show that fully private t5 models exhibit reasonable pre-training performance and don't hurt subsequent fine-tuning, and that private sentencepiece serves as a regularizer on noisy datasets and is able to improve pre-training and fine-tuning performance of models such as t5.