Plug and Play with Prompts:
A Prompt Tuning Approach for Controlling Text Generation

In our problem setting, we consider prompts as tunable soft-prompt embeddings (called prompt embeddings). Since our task is to generate completions for the opening phrase of a sentence, the opening phrase acts as the input text. This input text is converted to embeddings, referred to as source embeddings. The prompt embeddings are appended to the source embeddings as prefix, and fed to the transformer model as the input embeddings. Since our problem is concerned with open text generation, we use decoder models of the GPT-2 (Radford et al. 2019) family.

$$\mathbf{E_{i}}=\mathbf{[E_{p};E_{s}]}$$

(1)

where $\mathbf{E_{s}}$ are the source embeddings, $\mathbf{E_{p}}$ are the prompt embeddings, and $\mathbf{E_{i}}$ are the input embeddings.

Our framework for the generation task contains two models: Generator and Discriminator.

After a set number of generation steps, the generated text is concatenated to the source text and provided as the input to the discriminator (classifier) model. The discriminator produces the classification loss, which guides the prompts to generate text with a desired style. Backpropagating this loss to the prompt embeddings requires us to retain the gradient in the text generated by the generator. However, the tokens generated are discrete values obtained using argmax of the logits produced by the generator at each generation step, and they cannot preserve the gradient. To optimize the prompt through gradients, we approximate the argmax function using softmax with a low temperature (i.e., we divide the logits by a small number and take the softmax). This approximates the one-hot vector, while retaining the gradients. We then multiply this vector with the embedding matrices of the generator and classifier models to obtain the corresponding embedding vectors.

$$\mathbf{o_{1}}=G(\mathbf{E_{G_{i}}})$$

(2)

$$v_{1}^{i}=\frac{\exp(\frac{o_{1}^{i}}{\tau})}{\sum_{j\in V}\exp(\frac{o_{1}^{j}}{\tau})}$$

(3)

$$\mathbf{E_{G_{1}}}=\mathbf{v_{1}\cdot M_{G}}$$

(4)

The generator output at the first time step ($t$=$1$) is obtained by eqn. $2$. To approximate the one-hot vector ($\mathbf{v_{1}}$), we pass the output logits through a softmax function with a low temperature, as shown in eqn. $3$. $\mathbf{M_{G}}$ is the embedding matrix of the generator, which maps the vocabulary of the generator to the generator embeddings. The one-hot vector is multiplied with the generator embedding matrix to obtain the generator embedding at any time step $\mathbf{E_{G_{t}}}$.

Equations $2$ - $4$ show the generation of the generator embeddings at the first step ($t$=$1$). $\mathbf{o_{t}}$ is the logits vector generated at time $t$. $\mathbf{v_{t}}$ is the low temperature softmax of $\mathbf{o_{t}}$, used to approximate the one-hot vector. $\mathbf{M_{G}}$ is the embedding matrix of the generator, which maps the vocabulary of the generator to the generator embeddings. $\mathbf{E_{G_{t}}}$ and $\mathbf{E_{C_{t}}}$ refer to the generator and classifier embeddings generated at time $t$ respectively. $V$ is the vocabulary of the generator model.

The generator embedding is used to generate the next embedding at each generation step.

$$\mathbf{o_{t}}=G([\mathbf{E_{G_{i}}};\mathbf{E_{G_{1:t-1}}}])$$

(5)

Equations $3$ and $4$ hold true for any time step $t$, with $\mathbf{o_{1}}$ replaced by $\mathbf{o_{t}}$

After each generation step, the generated embedding is concatenated with the input embeddings for that generation step, and fed to the generator to obtain the output probabilities for the the next generation step.

$$\mathbf{E_{D_{t}}}=\mathbf{v_{t}\cdot M_{D}}$$

(6)

The discriminator embedding at any time-step $t$ are produced using eqn. $5$. $\mathbf{M_{D}}$ is the embedding matrix of the discriminator. Since the same one-hot vector $v_{t}$ is used to generate generator and discriminator embeddings, we require the generator and discriminator language models to have the same vocabulary.

To obtain the discriminator (classification) loss, input and output discriminator embeddings are concatenated and passed into the discriminator, as shown in eqn. $7$

$$\mathcal{L}_{D}=D([\mathbf{E_{D_{i}}};\mathbf{E_{D_{1:n}}}])$$

(7)

The prompts are expected to learn to produce text that can minimize the discriminator loss. This text should be both coherent and have the desired style. However, the prompts soon learn to produce text that fools the discriminator and minimizes the classification loss at the expense of coherence. In order to preserve the coherence, we use a second loss term, which acts as the fluency loss.

To set up our control, we use the same generator and the source text as the input, however, we do not attach the prompt embeddings. Thus, at each generation step of autoregressive generation, we obtain the logits that the generator would produce with the non-prompted source text (i.e., the original text). We call these the non-prompted logits, whereas the logits generated by the generator receiving the prompt and source embeddings are referred to as ‘prompted logits’. At each step of generation, we calculate the categorical cross-entropy between the prompted and non-prompted logits, and use the average over all generation steps as the fluency loss. The fluency loss is finally multiplied with a small number ($\lambda$), before being added to the discriminator loss for backpropogation.

We use two sentiment datasets, SST-5 (Socher et al. 2013) and Yelp reviews dataset from Shen et al. (2017), the GYAFC formality dataset (Rao and Tetreault 2018), and the Toxic Comment Classification Challenge (Jigsaw 2017) dataset to train the discriminator models. Table 2 further elaborates the datasets used to train the discriminator model.

To train the prompts, we create two types of datasets: in-domain and out-of-domain. Each entry in a dataset is an incomplete opening phrase, which is the input to the language model. The language model is expected to generate the completion for the opening phrase for a fixed number of generation steps.

The in-domain dataset is a very small set of inputs that closely resembles the training dataset of the discriminator model. For example, when using the discriminator trained on Yelp reviews, the in-domain dataset refers to restaurants, shops and hotels. Similarly, the in-domain dataset consists of conversation openings when using the discriminator trained on the GYAFC formality dataset. We use a $train:test$ split of $480:64$ data samples for all four style control problems.

On the other hand, the out-of-domain (OOD) dataset is a large collection of general sentence openings, consisting of a subject and a verb. The distribution of this dataset is independent of the datasets that are used to train the discriminators. This dataset has a $train:test$ split of $4800:320$ samples.

Both kinds of datasets (in-domain and out-of-domain) are created synthetically, using GPT4 (OpenAI 2023).

We perform automatic evaluation for both style and fluency using the following metrics:

1. 1.

   Style Accuracy: We evaluate the style of the generated text using an external classifier.

2. 2.

   Perplexity: We use perplexity as an automatic measure of fluency of the generated text, similar to Dathathri et al. (2020). Perplexity is calculated on the same model used for generating text (GPT2 Large).

3. 3.

   Dist-n: Dist-n is a measure of distinct n-grams in the generated text. We calculate the Dist-1, Dist-2 and Dist-3 scores for each method. A higher number of distinct n-grams is indicative of more diverse text (Li et al. 2016).

We conduct an ablation study with 5 variants:

1. 1.

   B: We provide the source text as input to the base model, and generate one sample.

2. 2.

   BR: Similar to B, we provide the source text as input to the base model (GPT2 Large), and generate $r$ samples. $r$ is set to 3 samples. We select the best sample according to perplexity and dist scores.

3. 3.

   BP: Here, we attach the prompt embeddings, tuned only using the discriminator loss, to the source embeddings, and use these as input to the base model.

4. 4.

   BPF: The prompt embeddings used here are tuned using both discriminator and fluency loss. These embeddings are attached to the source embeddings, and provided as input to the base model.

5. 5.

   BPFR: This is a sampled version of BPF, where the prompted model is sampled $r$ times, and the best sample chosen according to perplexity and dist scores. $r$ is set to 3 samples.

In all of the above ablations, we use GPT2 Large as the base model.

Table 1 shows the results for the proposed approach (PPP). We compare this with commonly used methods for instructing large language models (viz. Zero Shot and Few Shot). We also compare our method to existing plug-and-play methods, including PPLM (Dathathri et al. 2020) and GeDi (Krause et al. 2021). Our method substantially outperforms previous baselines, even though the prompts are trained using a very small dataset. Moreover, the gain in performance compared to zero-shot and few-shot methods highlights how prompt-tuning is superior to using textual prompts, particularly in the case of smaller models. We note a slight trade-off between fluency (perplexity) and style control (style accuracy). However, this can be controlled using the fluency loss hyperparameter $\lambda$. Perplexity skyrockets in the case of approach BP, where the prompt embeddings are tuned using only the discriminator loss (i.e., $\lambda=0$). This highlights the need for using the fluency loss, which is calculated using Cross-Entropy between the prompted and unprompted language model. While GeDi slightly outperforms approach BPF (i.e., PPP without sampling), and PPP on the toxicity dataset, it is important to note that GeDi’s outputs show significantly higher perplexity than PPP.

Overall, our results show that plug-and-play methods based on prompt tuning have the potential for data-efficient controlled generation, with a high degree of fluency, especially when using smaller models.

In this section, we demonstrate the efficacy of PPP on a larger and more general, out-of-domain training dataset. Unlike the small datasets used for training the prompts previously, the dataset used here does not necessarily follow the same distribution as the style dataset used for training the discriminator model. We tune the prompt embeddings on the large dataset for two problems - sentiment control and formality control. Since we use a much larger training dataset, we also increase the number of trainable parameters (i.e., the number of prompt embeddings). Towards this, we train the prompts for GPT2 Large by varying the number of prompts in $\{25,50,75,100,125,150\}$ and observe the change in performance (style accuracy%) in figure 2a. We note that the performance steadily increases as the number of trainable prompt embeddings is increased. In addition to this, we also experiment with various model sizes: GPT2 ($117$ M parameters), GPT2 Medium ($345$ M parameters), GPT2 Large ($774$ M parameters), and GPT2 XL ($1.558$ B parameters). We note a similar trend in figure 2b, i.e., the performance increases with increasing the number of model parameters for both tasks.

In this section, we demonstrate the control and generation quality of our method by analyzing the text generated using the prompts trained using different classifiers.

Table 3 illustrates that the prompt embeddings are able to change the formality of the generated output, for the same input text, depending on the direction (formal or informal) that they were trained to produce.

Table 4 shows how the prompts exhibit greater control as they are trained over epochs. The output text is negative at first, and gradually changes to more positive, as training progresses.

Table 5 shows a comparison between our method and existing plug and play methods. Our method demonstrates higher fidelity towards the dataset used for training the discriminator, as it generates a very negative yelp review for a cafe.

Table 6 demonstrates an application of our method towards detoxifying text generated by a language model. The original text generated by GPT2 contains abusive text and profanities. Upon addition of prompt embeddings (trained to generate non-toxic text), the generated text does not show any toxic characteristics.

We tune the hyperparameters on SST-5 with GPT2 Large, and use the same hyperparameters for all 4 style control problems. We experiment with prompt lengths in $\{10,20,30,40,50,80\}$. We obtain the best results with prompt length of 30. We search the learning rate in $\{1e-5,3e-5,1e-4,3e-4,1e-3,3e-3,1e-2,3e-2\}$, and obtain the best results with $1e-3$. Higher learning rates lead to poor text quality and lower learning rates do not perturb the style significantly. We search the fluency loss parameter ($\lambda$) in $\{0.05,0.1,0.15,0.2,0.25,0.3,0.4,0.5\}$, and obtain the best results with $\lambda=0.15$. For all cases, we use the AdamW optimizer (Loshchilov and Hutter 2019). We train for a maximum of $20$ epochs on the small datasets, and $5$ epochs on the large OOD dataset.