InfoPrompt: Information-Theoretic Soft Prompt Tuning for Natural Language Understanding

Soft prompt tuning has shown great successes in a wide range of NLP tasks [42, 24, 43]. Let $\Phi$ denote the encoder of a pretarined language model, e.g., Roberta-Large [47]. Assume $X=\{x_{1},x_{2},\cdots,x_{n}\}$ is a text sequence of length $n$ and $Y$ is its label. By prompting, we add extra information $P$ for the model to condition on during its prediction of $Y$. $P=\{p_{1},\dots,p_{n_{p}}\}$ is a sequence of token embeddings. $n_{p}$ is the number of prompt token. $p_{i}\in\mathbb{R}^{D}$, $i=1,\cdots,n_{p}$, is a embedding vector with dimension $D$ and $D$ is the embedding dimension of the pretrained language model. We first embed each token of $X$ into its corresponding token embedding from the pretrained language model. $P$ is inserted into the the resulting embedding sequence of $X$ to form a template, which is further encoded by the pretrained encoder $\Phi$ into the representation space of the pretrained langauge model. The template is detailed with the experiments. Formally, we denote such a process as $Z=\Phi(P,X)$, with $Z$ being the encoded representation. The model prediction for $X$ is made on top of $Z$ with a trainable language modeling head $h_{\theta}$ (parameterized by $\theta$), s.t., the output $h_{\theta}(Z)$ is the distribution over all possible labels for classification. The whole model is generally trained via minimizing the following loss function,

Mutual information (MI) is a metric in information theory [60, 14], which measures the amount of shared information between two random variables. The mutual information between two random variables $A$ and $B$ is

Inspired by [63], we use MI as the criterion for comparing prompts and other model parameters. MI has also been applied to measure the similarity between the masked token and the corresponding context in the pretraining of Multilingual Masked Language Modeling [9], the relevance between the document and a sentence in document summarization [55] and the source sentence and the target translated sentence in Neural Machine Translation (NMT) [80].

the head loss is the negative mutual information between the prompt $P$ and parameters $\theta$, i.e., $-I(P;\theta|X)$. In maximizing $I(P;\theta|X)$, we follow [52] that approximate it with the following lower bound,

We model the score function $l(\theta,p)$ as a standard bilinear function with the learnable matrix $W_{1}$

where $\theta$ and $P$ are encoded from $X$. $W_{1}$ is a trainable matrix. Since the language model head is learnt on top of the output from the last layer of the pretrained language model, the learning of its parameters $\theta$ is easier than the learning of the prompt $P$ (the input layer of the pretrained language model). Therefore, $\theta$ may capture more task-relevant information than $P$ in the early stage of training. By maximizing the mutual information between $\theta$ and $P$, the task-relevant information can be transferred to $P$ in the initial training steps. In this way, $P$ can be more task-relevant especially in the early training stage. Experiments in Section 6 also show that our head loss, $\mathcal{I}(P;\theta|X)$, can facilitate the training of the initial training steps.

We minimize the following objective in prompt tuning:

We denote $\mathcal{L}_{\mathrm{pred}}$ as the task loss. $\beta$ and $\gamma$ are balancing parameters for the proposed representation loss and head loss, respectively. We denote our approach as InfoPrompt.

We state our main theoretical result as follows. Due to the space limit, we delay the proof into Appendix.

We provide theoretical guarantees of the convergence of those losses trained by conventional gradient decent type algorithms, which enables our method to converge more easily than traditional soft prompt tuning methods and avoid the problem of large training epochs.

We conduct experiments with datasets of sequence classification from the GLUE benchmark [74], along with those of relation extraction tasks and NER tasks. We choose four seqience classification tasks from the GLUE benchmark: RTE (Recognizing Textual Entailment, [4]), MRPC (Microsoft Research Paraphrase Corpus, [16]), CoLA (Corpus of Linguistic Acceptability, [73]) and SST-2 (Sentence Sentiment Treebank, [61]). We choose these tasks because their datasets are in smaller sizes and prompt tuning is comparably more effective in low-resource setting [33, 45]. For the task of relation extraction, we evaluate our method on ACE2005 corpus and Semeval-2010 datasets [32]. We also use the ACE2005 corpus for the task of NER. Note that the entity spans for NER have been given ACE2005. Unlike the standard NER model that learns to predict the entity span and entity type simultaneously from the raw text sequence, our model only predicts the entity type based on the given entity span. We follow the same data splitting strategy for ACE2005 corpus as the previous work [78, 54]. For the Semeval-2010 tasks, we follow the official data partition [32].

We follow the resource constrained scenario in [33] that train each task with only 64 or 256 samples. We experiment with $n_{p}=1$ and $n_{p}=4$ prompt tokens for each task. The prompt tokens are inserted into the template for each task. Similar to [33], we adopt the RoBERTa-large model as our pretrained encoder. We freeze pretrained parameters are only trained with parameters of the prompt head and prompt tokens. During training, we empirically set $\beta=0.1$ and $\gamma=0.05$. The number of negative of negative samples $K=32$. The learning rate is $1e-3$ and the batch size is 8. For each task, we report the results after 30 epoches, averaged over 5 random seeds. To encode the prompt $P=[p_{1},\cdots,p_{n_{p}}]$ from $X$, we first encode $X$ into $P^{\prime}\in\mathbb{R}^{D}$ via $P^{\prime}=\Phi(X)$. For each $p_{i}\in\mathbb{R}^{D}$, we have $p_{i}=W_{i}^{\mathrm{up}}W_{i}^{\mathrm{down}}P^{\prime}$, $W_{i}^{\mathrm{up}}\in\mathbb{R}^{D\times 64}$, $W_{i}^{\mathrm{down}}\in\mathbb{R}^{64\times D}$.

For the tasks of sequence classification and relation extraction, we follow the template of [33] that contains a $\mathsf{[mask]}$ token. The representation $Z$ is the $\mathsf{[mask]}$ token from the last layer of the Roberta-Large encoder. For the tasks of named entity recognition, we have the $\mathsf{[mask]}$ token before the given entity span, with the rest being the same as for sequence classification.

As mentioned above, our method with (11) in denoted as InfoPrompt. In the experiments, we compare our method with the following baselines:

• •

  Fine Tuning: We fine tune all the parameters from the pretrained encoder on each task. Fine Tuning is included as the upper bound for the model performance, since it is more computational expensive compared with only training the prompt parameters.

• •

  Adapter: Similar to prompt tuning, this is also a way of parameter-efficient training for pretrained language models. Specifically, instead of adding the prompt tokens in the input, we add adaptors after the last linear layer in each transformer layer.

• •

  WARP [33]: Different from our approach, the prompt tokens of WARP are not generated from the input sequence. the prompt tokens are insert into the input sequence. During training, the pretrained encoder is frozen and only the prompt tokens are trainable.

• •

  IDPG [76]: Similar to our approach, the prompt tokens are generated from the input sequence. The pretrained encoder is frozen and the prompt generator is trainable.

In evaluating the effectiveness of our proposed loss functions, we consider the following two ablations:

• •

  $\gamma=0$: We disable the head loss during training via $\gamma=0$, while keeping $\beta=0.05$.

• •

  $\beta=0$: We disable the head loss during training via $\gamma=0$, while keeping $\gamma=0.1$.

• •

  $\beta=\gamma=0$: We disable both the losses. The prompt parameters are trained with $L_{\mathrm{pred}}$.

Table 1 and 2 show that results of training with the full dataset for each task. We can observe that the results with our InfoPrompt is generally higher than the other parameters-efficient baselines that freeze the pretrained Roberta-Large parameters (e.g., WARP and Adapter). Fine Tuning generally have better performance than the other approaches. This is because it allows training with all the model parameters, which is at the expense of more computation cost during training. As mentioned above Fine Tuning is intended to be included as the upper bound for performance. Moreover, we can find the performance with $\gamma=0$ and $\beta=0$ is lower than that of InfoPrompt, indicating shows that it is beneficial to learn task-relevant prompt tokens with the propose head loss and representation loss. Further, the performance gap between $\gamma=0$/$\beta=0$ and $\beta=\gamma=0$ shows that the proposed functions are effective when added to naive prompt tuning, i.e., with only $\mathcal{L}_{pred}$.

The results for training with few-shot datasets are listed in Table 3 and 4. Compared with training with the full dataset (Table 1 and 2), we can find that the performance gap between our propose InfoPrompt and the baselines is generally larger in the few-shot setting. Unlike the full datasets, the few-shot datasets contains much less information regarding the task to be learnt. As the result, the prompts learnt with solely the task loss (e.g., with WARP or $\beta=\gamma=0$) may easily overfit to the task-irrelvant information given few-shot datasets. In such a scenario, it would be important to explicitly encourage the learnt prompt to be task-relevant, i.e., via our proposed loss functions based on mutual information maximization. This explains why InfoPrompt yields larger performance gain when trained with few-shot datasets. Similar to training with the full datasets, the performance gains of InfoPrompt compared with InfoPrompt ($\gamma=0$) and InfoPrompt ($\beta=0$) show the effectuveness of our proposed loss functions in the few-shot scenario.

To provide a more comprehensive understanding on the effectiveness of our proposed loss functions, we plot the landscapes of the loss functions in the parameter space of the prompt tokens. The landscapes illustrates how the values of loss functions varies with the input prompt tokens. Since the prompt tokens are high-dimensional vertors, i.e., each token has the dimension of 1024 for RoBerta-Large, we visualize their associated loss values via projecting the prompt tokens into a 2D subspace. Specifically, we follow previous work of token embedding analysis [10] that project the prompt tokens into the top-2 principled component computed from the pretrained token embeddings of Roberta-Large. We only insert one prompt token into the input sequence during visualization.

Taking the task of MRPC as an example, we plot the 2D landscapes of the task loss and the representation loss in Figure 2(a) and 2(b), respectively. All the two figures are ploted with the same scale, i.e., with the same values of the prompt token. The axis values are the offset from the mean of the pretrained Roberta-Large token embeddings. The loss values shown in the figures are the average of 20 random samples from MRPC. In Figure 2(a), we can find that there are a lot of local minimum in the landscapes of the task loss. This is consistent with the observations of the previous works [24, 65] that prompt tuning is difficult to be optimized with and sensitive to initialization, e.g., the optimization can get easily overfit to a local minimum without proper initialization. From Figure 2(b), we can observe that the values of our proposed info loss is much smoother compared with the task loss in Figure 2(a). With smoother landscapes, the optimization with our proposed loss functions can be more stable (also shown in Section 6.2), i.e.. less likely being trapped in a local minimum and also guaranteed to converge by our theoretical results (see Theorem 3.1). Additionally, we plot the trajectory of the first 500 steps during training for InfoPrompt ($\gamma=0.05$) (green) and $\gamma=0$ (purple) in Figure 2(a) and 2(b). The stars in the plot indicate the initial value of the prompt before training. We find that training with $\gamma=0.05$ can render a larger descent for both the task loss and representation loss, compared with $\gamma=0$. As analyse in Section 3.1, the language head is easier to be learnt than the prompt. As the result, parameters of the language head may contain more task-relevant information during the earlier stage of training. By maximization the mutual information between the head parameter and prompt via the proposed head loss (weighted by $\gamma$), we encourage the learnt prompt to capture more task-relevant information in the initial training steps, thus resulting $\gamma=0.05$ to have a larger descent than $\gamma=0$ in the trajectories shown in Figure 2(a) and 2(b). Note that our propose two loss functions are unsupervised and do not required additional labels.

We plot the training curve for the task of NER and SST-2 in Figure 3(a) and 3(b), respectively. Unlike WARP [33] and Fine Tuning that train with solely the task loss $L_{\mathrm{pred}}$, our InfoPrompt also train with the representation loss and head loss. We can observe that the training of our our InfoPrompt is more stabilized and converges faster, compared with WARP. This can be explained from the landscape plots in Section 6.1. Since the landscape of the task loss is not smooth (Figure 2(a)), the training curve of WARP may exhibit significant perturbation when the optimization overfits to a local minimum, e.g., the 10000th in Figure 3(a). Comparably, our proposed InfoPrompt can smooth the optimization landscape, thus stabilize the training and result in faster convergence, which is guaranteed by our theoretical results. We observe that Fine Tuning generally converges faster and end up with a higher accuracy than InfoPrompt. This is because Fine Tuning, which trains with all the model parameters, have much larger model capacity during training than prompt tuning (InfoPrompt and WARP). Such results for Fine Tuning is at the expense of larger computation cost, i.e., we need to calculate the gradient for all the model parameters (354M) instead of only the prompt parameters $P$ (1.3M).

Soft prompt tuning becomes a new paradigm in NLP. Based on some large pretrained models (e.g., BERT [17], RoBERTa [47]), a relatively small number of trainable parameters can be added to the input, while the parameters of backbones are fixed. Many works have demonstrated the effectiveness of soft prompt tuning in a wide range of NLP downstream tasks [42, 33, 56], especially in low-resource regions [62, 43, 24]. Some recent works also found the transferable power of soft prompts across domains [72, 65, 68], across language models [65] and for zero-shot generalization [77]. To further enhance the efficiency of soft prompt parameters and enable better generalization abilities, many works consider multi-task learning [2, 23, 68, 36], or multilingual [11, 34]. Some works also try to explore the prompt with prior knowledge encoded [30, 35, 15]. While most of the initial attempts of soft prompt tuning are not context-aware, some recent works suggest that the soft prompt tokens should be conditioned on the input context. Hyperprompt [36] proposes a hyper-network structure to generate prompt tokens based on task indexes. [76] and [7] suggest some context-aware soft prompt generation methods. [46] proposes a structured soft prompt tuning method.

Information-theoretic methods are widely used in many NLP tasks [38, 69, 59, 40, 53]. [69] and [38] propose information-theoretic methods for text memorization. [53] suggests an information-theoretic method for dialogue. [40] views the multimodal NMT problem in an information-theoretic point of view. For model pretraining, [75] proposes Infobert to improve the robustness of the BERT model. INFOXLM [9] proposes a cross-lingual language model based on an information-theoretic framework. For model fine-tuning, [51] proposes an information bottleneck model method for low-resource fine-tuning. [63] proposes an information-theoretic method to engineer discreet prompts.

Softmax is one of the major unit in the attention scheme of most recent NLP large language models. Computing the attention matrix faster is a practically interesting question [12, 71, 41]. Recently, a number of theoretical work have tried to study the softmax/attention unit from theoretical perspective. The softmax attention computation can be formally defined as follows: suppose we are given three matrices $Q\in\mathbb{R}^{n\times k}$ (the query), $K\in\mathbb{R}^{n\times k}$ (the key), and $V\in\mathbb{R}^{n\times k}$ (the value), the goal is to compute $\mathsf{Att}(Q,K,V)=D^{-1}\exp(QK^{\top})V$ where the diagonal matrix $D$ is $\mathrm{diag}(\exp(QK^{\top}){\bf 1}_{n})$. Here $K^{\top}$ denote the transpose of matrix $K$. The work of [79, 1] consider the static setting, and the work of [8] considers the dynamic setting. [1] proposed a tight algorithm for computing $\mathsf{Att}$ and provided a lower bound result based on the strong exponential time hypothesis. The work [8] shows a tight positive result and a negative result. In [8], they provide an upper bound via lazy update techniques [13]. In [8], they also present a lower bound result which is based on the Hinted MV conjecture [5]. The work of [22] proposes two sparsification algorithm to compute attention matrix when the feature dimension $\gg$ the length of sentence. [27] shows how to provide a differentially private algorithm for computing attention matrix under differential privacy framework [21, 18].