SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1 and 2 show the results of training with the full dataset for each task. We can observe that the results with our InfoPrompt are generally higher than those of the other parameters-efficient baselines that freeze the pretrained RoBERTa-Large parameters (e.g., WARP and Adapter). Finetuning generally has 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 Finetuning is intended to be included as the upper bound for performance. Moreover, we can find that 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 proposed 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}$ .
5.2 Training with the Few-Shot Datasets
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 proposed InfoPrompt and the baselines is generally larger in the few-shot setting. Unlike the full datasets, the few-shot datasets contain 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-irrelevant information given the 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 effectiveness of our proposed loss functions in the few-shot scenario.
(a) The landscape of task loss.  $\beta = 0.1$
(b) The landscape of representation loss.  $\beta = 0.1$
Figure 2: 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. The trajectory shows the first 500 steps during training for InfoPrompt with  $\gamma = 0.05$  or  $\gamma = 0$ .
6.1 Loss Landscapes
To provide a more comprehensive understanding of 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 illustrate how the values of the loss functions vary with the input prompt tokens. Since the prompt tokens are high-dimensional vectors, 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 on token embedding analysis [12] that projects the prompt tokens into the top-2 principal components 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 2a and 2b, respectively. Both figures are plotted 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 2a, 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 [37, 91] 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 2b, we can observe that the loss landscape of our proposed representation loss is much smoother compared to the task loss in Figure 2a. With smoother landscapes, the optimization with our proposed loss functions can be more stable (also shown in Section 6.2), i.e., less likely to be trapped in a local minimum and also guaranteed to converge according to 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 2a and 2b. 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 to  $\gamma = 0$ . As analyzed in Section 3.1, the language head is easier to learn than the prompt. As the result, parameters of the language head may contain more task-relevant information during the earlier stage of training. By maximizing 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 2a and 2b. We also compare our initialization to some common initialization approaches: Random Uniform and Sampled Vocabulary [60, 37]. By Random Uniform, we randomly sample prompt initialization from the continuous latent space. By Sampled Vocabulary, we randomly sample prompt initialization from language model's vocabulary set. The final results by WARP (Random Uniform), WARP (Sampled Vocabulary) and InfoPrompt are 0.626, 0.672 and 0.706 respectively. The results
(a) NER learning curve.
(b) SST-2 learning curve.
Figure 3: The learning curves for the task of NER and SST-2. The training of our our InfoPrompt is more stabilized and converges faster, compared with WARP.
validate the effectiveness of our initialization approach. Note that our proposed two loss functions are unsupervised and do not require additional labels.
6.2 Learning Curve
We plot the training curve for the task of NER and SST-2 in Figure 3a and 3b, respectively. Unlike WARP [40] and Finetuning that train with solely the task loss  $L_{\mathrm{pred}}$ , our InfoPrompt also trains with the representation loss and head loss. We can observe that the training of 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 2a), the training curve of WARP may exhibit significant perturbation when the optimization overfits to a local minimum, e.g., the 10000th step in Figure 3a. Comparably, our proposed InfoPrompt can smooth the optimization landscape, thus stabilizing the training and result in faster convergence, which is guaranteed by our theoretical results. We observe that Finetuning generally converges faster and ends up with a higher accuracy than InfoPrompt. This is because Finetuning, which trains with all the model parameters, has much larger model capacity during training than prompt tuning (InfoPrompt and WARP). Such results for Finetuning is at the expense of larger computation costs, i.e., we need to calculate the gradient for all the model parameters (354M) instead of only the prompt parameters  $P$  (1.3M).
We also validate that our approach is less sensitive to initialization in the early learning stage, compared to WARP. Specifically, across 10 different random seeds, we report the standard errors of the performances by InfoPrompt and WARP in the early learning stage. After the first epoch, on NER, the results by WARP and InfoPrompt are  $0.461 \pm 0.038$  and  $0.810 \pm 0.025$ , respectively. On SST2, the results by WARP and InfoPrompt are  $0.735 \pm 0.033$  and  $0.764 \pm 0.027$ , respectively. The results show that our method has lower standard errors and is less sensitive compared to WARP.
7.1 Soft Prompt Tuning
Soft prompt tuning has become a new paradigm in NLP. Based on some large pretrained models (e.g., BERT [25], RoBERTa [67]), 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 [60, 40, 73, 65], especially in low-resource regions [78, 66, 37]. Some recent works also found the transferable power of soft prompts across domains [101, 91, 94], across language models [91] and for zero-shot generalization [105]. To further enhance the efficiency of soft prompt parameters and enable better generalization abilities, many works consider multi-task learning [6, 27, 94, 43], or multilingual [14, 50]. Some works also try to explore the prompt with prior knowledge encoded [48, 42, 13]. 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 [43] proposes a hyper-network structure to generate prompt tokens based on task indexes. [102] and [8] suggest some context-aware soft prompt generation methods. [64] proposes a structured soft prompt tuning method. BBT [92] targets the scenarios where the pre-trained model is not available locally (i.e., deployed online) and its back-propagation operation is not available.
7.2 Information-theoretic Methods in NLP
Information-theoretic methods are widely used in many NLP tasks [55, 99, 90, 51, 70]. [99] and [55] propose information-theoretic methods for text memorization. [70] suggests an information-theoretic method for dialogue. [51] views the multimodal NMT problem in an information-theoretic point of view. For model pretraining, [96] proposes Infobert to improve the robustness of the BERT model. INFOXLM [15] proposes a cross-lingual language model based on an information-theoretic framework. For fine-tuning, [69] proposes an information bottleneck model method for low-resource fine-tuning. [89] introduces an information-theoretic method to engineer discrete prompts.
7.3 Theoretical Attention Computation
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 [17, 97, 57]. 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  $\mathrm{Att}(Q, K, V) = D^{-1} \exp(QK^{\top}) V$  where the diagonal matrix  $D$  is  $\mathrm{diag}(\exp(QK^{\top}) \mathbf{1}_n)$ . Here  $K^{\top}$  denote the transpose of matrix  $K$ . The work of [106, 3] consider the static setting, and the work of [11] considers the dynamic setting. [3] proposed a tight algorithm for computing Att and provided a lower bound result based on the strong exponential time hypothesis. [4] provide the results for a more general tensor version of attention which capture the three tuples feature, but classical attention cannot [77]. The work [11] shows a tight positive result and a negative result. In [11], they provide an upper bound via lazy update techniques [18]. In [11], they also present a lower bound result which is based on the Hinted MV conjecture [10]. The work of [21] proposes two sparsification algorithms to compute attention matrix when the feature dimension  $\gg$  the length of sentence. [35] shows how to provide a differentially private algorithm for computing attention matrix under differential privacy framework [30, 29]. [41] introduces a hyperattention method and presents an nearly linear time algorithm with provable guarantees. [56] studies the polynomial based attention scheme and shows that sketching techniques can help speeding up the attention computation.
We revisit the limitations of soft prompt tuning in the initialization. We also empirically discover that conventional prompt tuning methods cannot learn sufficient task-relevant information from prompt tokens. We tackle these limitations from an information-theoretic perspective and propose an information-theoretic prompt tuning method InfoPrompt, with two novel loss functions. With extensive experiments, without any prior expert knowledge, InfoPrompt can significantly accelerate the convergence of the prompt tuning and achieve more accurate and robust performances than traditional prompt tuning methods.
Existing instruction-tuned LMs (e.g., Llama 2 [93]) are generally not task-specific and future works may further consider to tune such models with task-specific information. To achieve this, combining soft prompt tuning and prompt engineering [104, 76], from an information-theoretical perspective, can be a promising approach. Another future direction could be to further generalize our method to generation tasks (e.g., sequence generation). In addition to prompt learning, it is interesting to explore how to extend our approach to other parameter-efficient fine-tuning methods (e.g., LoRA [47] and HyperFormer [26]).
The authors would like to thank the anonymous reviewers for their insightful comments. During this research, Ricardo Henao and Rui Wang were supported by ONR N00014-18-1-2871-P00002-3.
[102] Zhuofeng Wu, Sinong Wang, Jiatao Gu, Rui Hou, Yuxiao Dong, VG Vydiswaran, and Hao Ma. Idpg: An instance-dependent prompt generation method. arXiv preprint arXiv:2204.04497, 2022.
[103] Bishan Yang and Tom Mitchell. Joint extraction of events and entities within a document context. arXiv preprint arXiv:1609.03632, 2016.
[104] Kexin Yang, Dayiheng Liu, Wenqiang Lei, Baosong Yang, Mingfeng Xue, Boxing Chen, and Jun Xie. Tailor: A soft-prompt-based approach to attribute-based controlled text generation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 410–427, 2023.
[105] Seonghyeon Ye, Joel Jang, Doyoung Kim, Yongrae Jo, and Minjoon Seo. Retrieval of soft prompt enhances zero-shot task generalization. arXiv preprint arXiv:2210.03029, 2022.
[106] Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In ICML. arXiv preprint arXiv:2302.02451, 2023.
[107] Songming Zhang, Yijin Liu, Fandong Meng, Yufeng Chen, Jinan Xu, Jian Liu, and Jie Zhou. Conditional bilingual mutual information based adaptive training for neural machine translation. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2377–2389, 2022.
[108] Kai Zhong, Zhao Song, Prateek Jain, Peter L Bartlett, and Inderjit S Dhillon. Recovery guarantees for one-hidden-layer neural networks. In International conference on machine learning (ICML), pages 4140-4149. PMLR, 2017.
Roadmap. In Section A, we provide a number of basic notations. In Section B, we provide several basic definitions and discuss some related work about previous theoretical softmax regression results. In Section C, we provide a complete proof for our major theoretical result in this paper. We present our final result in Section D.
For any positive integer  $n$ , we use  $[n]$  to denote set  $\{1, 2, \dots, n\}$ . For any function  $f$ , we use  $\widetilde{O}(g)$  to denote  $g \cdot \mathrm{poly}(\log g)$ .
Vector For a length- $n$  vector  $z$ , we use  $\exp(z)$  to denote a length- $n$  vector that its  $i$ -th entry is  $\exp(z_i)$ .
For a length- $n$  vector  $z$ , we use  $\|z\|_2$  to represent its  $\ell_2$  norm, i.e.,  $\|z\|_2 := (\sum_{i=1}^n x_i^2)^{1/2}$ . For a length- $n$  vector  $z$ , we use  $\|z\|_\infty$  to denote  $\max_{i \in [n]} |z_i|$ .
For a length- $n$  vector  $z \in \mathbb{R}^n$ , we use  $\mathrm{diag}(z)$  to generate a diagonal matrix where each entry on the  $(i, i)$ -diagonal is  $z_i$  for every  $i \in [n]$ .
We use  $\mathbf{1}_n$  to represent a length- $n$  vector where all the coordinates are ones. Similarly, we use  $\mathbf{0}_n$  to represent a length- $n$  vector where all the values are zeros.
PSD We say  $W \succeq Z$  (positive semi-definite) if  $x^{\top}Wx \geq x^{\top}Zx$  for all vector  $x$ .
Matrix Related For an  $n$  by  $d$  matrix  $C$ , we use  $\mathrm{nnz}(C)$  to denote the number of non-zero entries of  $C$ , i.e.,  $\mathrm{nnz}(C) := |\{(i,j) \in [n] \times [d] \mid C_{i,j} \neq 0\}|$
For a diagonal matrix  $D \in \mathbb{R}^{n \times n}$ , we say  $D$  is a  $k$ -sparse diagonal matrix, i.e.,  $k = |\{i \in [n] \mid D_{i,i} \neq 0\}|$ .
For any matrix  $Z \in \mathbb{R}^{n \times k}$ , we denote the spectral norm of  $Z$  by  $\| Z\|$ , i.e.,
$$
\| Z \| := \sup  _ {x \in \mathbb {R} ^ {k}} \frac {\| Z x \| _ {2}}{\| x \| _ {2}}.
$$
For a matrix  $Q$ , we use  $\sigma_{\max}(Q)$  to denote the largest singular value of  $Q$ . We use  $\sigma_{\min}(Q)$  to denote the smallest singular value of  $Q$ .
Matrix Computation We use  $\omega$  to denote the exponent of matrix multiplication, i.e.,  $n^{\omega}$  denotes the time of multiplying an  $n\times n$  matrix with another  $n\times n$  matrix. Currently  $\omega \approx 2.373$  [100, 58, 5].
Calculus Related We use  $\circ$  notation by following the literature's [63, 20, 36, 61, 80]. Suppose that we're given two column vectors  $x, y \in \mathbb{R}^n$ , we use  $x \circ y$  to denote a column vector that  $(x \circ y)_i$  is  $x_i y_i$ .
B Related Work about Theoretical Attention Regression Results
In this paper, we focus on the direction of regression tasks [32, 63, 20, 61, 80, 34, 23, 16, 33, 82, 24]. The goal of this section will review the linear regression (Definition B.1), exponential regression (Definition B.2), rescaled softmax regression (Definition B.3), softmax regression (Definition B.2).
Definition B.1 (Linear regression). Given a matrix  $A \in \mathbb{R}^{n \times d}$  and  $b \in \mathbb{R}^n$ , the goal is to solve
$$
\min  _ {x \in \mathbb {R} ^ {d}} \| A x - b \| _ {2}.
$$
For convenient, let us  $u(x)$  to denote  $\exp (Ax)$ .
Definition B.2 (Exponential Regression, see [32, 63]). Suppose we are given a length  $n$  vector  $b$ , and an  $n$  by  $d$  size matrix  $A$ , our goal is to optimize
$$
\min  _ {x \in \mathbb {R} ^ {d}} \| u (x) - b \| _ {2}.
$$
Definition B.3 (Rescaled Softmax Regression, see [36]). Suppose we are given a length  $n$  vector  $b$ , and an  $n$  by  $d$  size matrix  $A$ , our goal is to optimize
$$
\min  _ {x \in \mathbb {R} ^ {d}} \| u (x) - \langle u (x), \mathbf {1} _ {n} \rangle \cdot b \| _ {2}
$$
Definition B.4 (Softmax Regression, see [20, 61, 80]). Suppose we are given a length  $n$  vector  $b$ , and an  $n$  by  $d$  size matrix  $A$ , our goal is to optimize
$$
\min  _ {x \in \mathbb {R} ^ {d}} \| \langle u (x), \mathbf {1} _ {n} \rangle^ {- 1} \cdot u (x) - b \| _ {2}.
$$
C Theoretical Guarantees
In Section C.1, we provide several basic definitions. In Section C.2, we explain how to compute the gradient of function  $f$ . In Section C.3, we show how to compute the gradient of function  $\log f(x)$ . In Section C.4, we explain how to compute the Hessian of function  $\log f(x)$ . In Section C.5, we compute the hessian of inner product between  $\log f(x)$  and  $b$ . In Section C.6, we compute the Hessian of cross entropy loss function. In Section C.7, we show Hessian is positive definite. In Section C.8, we prove that Hessian is Lipschitz. We remark that our experiments are based on first order method, and our theoretical proofs are mainly focusing on second order method. We believe it's an interesting future direction to further study the convergence of the first order method such as [108, 1, 2, 88, 84, 83].
C.1 Function Definition
Definition C.1. We define  $u(x)$  as follows
$$
\bullet u (x) = \exp (A x)
$$
Definition C.2. We define  $v(x)$  as follows
$$
\bullet v (x) = \exp (A x)
$$
Previous [63] studies three types of hyperbolic functions  $\exp (\cdot)$ ,  $\cosh (\cdot)$  and  $\sinh (\cdot)$ . We mainly focus on  $\exp (\cdot)$  function.
Definition C.3 (Normalized coefficients, Definition 5.4 in [20]). We define  $\alpha : \mathbb{R}^d \to \mathbb{R}$  as follows
$$
\alpha (x) := \langle u (x), \mathbf {1} _ {n} \rangle .
$$
We define function softmax  $f$  as follows
Definition C.4 (Function  $f$ , Definition 5.1 in [20]). Suppose that we're given an  $n \times d$  matrix  $A$ . Let  $\mathbf{1}_n$  denote a length- $n$  vector that all the coordinates are ones. We define prediction function  $f: \mathbb{R}^d \to \mathbb{R}^n$  as follows
$$
f (x) := \langle u (x), \mathbf {1} _ {n} \rangle^ {- 1} \cdot u (x).
$$
Fact C.5. Let  $f(x)$  be defined as Definition C.4. Then we have
- Part 1.  $\langle f(x),\mathbf{1}_n\rangle = 1$
- Part 2.  $\| f(x)\| _1 = 1$
Part 3.  $\| f(x)\| _2\leq 1$
Proof. The proof is straightforward. For more details, we refer the readers to [20].
We define the  $\ell_2$  loss
Definition C.6. We define
$$
L _ {\exp} := 0. 5 \| f (x) - b \| _ {2} ^ {2}.
$$
Previous work [80] only considers entropy, here we consider cross entropy instead.
Definition C.7 (Cross Entropy). We define  $L_{\mathrm{cent}}:\mathbb{R}^d\to \mathbb{R}$
$$
L _ {\text {c e n t}} (x) := - \langle b, \log f (x) \rangle
$$
Definition C.8. Suppose we're given an  $n \times d$  matrix  $A$  and  $W = \mathrm{diag}(w) \in \mathbb{R}^{n \times n}$  where  $w \in \mathbb{R}^n$  is a vector, we define  $L_{\mathrm{reg}}: \mathbb{R}^d \to \mathbb{R}$
$$
L _ {\mathrm {r e g}} (x) := 0. 5 \| W A x \| _ {2} ^ {2}
$$
C.2 Gradient Computation for Function  $f$
We present a calculus tool from previous work [20] (for example, we refer the readers to Lemma 5.6 in [20]).
Lemma C.9. If the following conditions hold
- Given matrix  $A \in \mathbb{R}^{n \times d}$  and a vector  $b \in \mathbb{R}^n$ .
- Suppose that function  $\alpha : \mathbb{R}^d \to \mathbb{R}$  be defined in Definition C.3.
- Suppose that function  $f: \mathbb{R}^d \to \mathbb{R}^n$  be defined in Definition C.4.
For each  $i\in [d]$  , we have
$$
\frac {\mathrm {d} f (x)}{\mathrm {d} x _ {i}} = - \langle f (x), A _ {* , i} \rangle \cdot f (x) + f (x) \circ A _ {* i}
$$
$$
\langle \frac {\mathrm {d} f (x)}{\mathrm {d} x _ {i}}, A _ {*, i} \rangle = - \langle f (x), A _ {*, i} \rangle^ {2} + \langle f (x), A _ {*, i} \circ A _ {*, i} \rangle
$$
$$
\langle \frac {\mathrm {d} f (x)}{\mathrm {d} x _ {i}}, A _ {*, j} \rangle = - \langle f (x), A _ {*, i} \rangle \cdot \langle f (x), A _ {*, j} \rangle + \langle f (x), A _ {*, i} \circ A _ {*, j} \rangle
$$
C.3 Gradient Computation for Function  $\log f(x)$
In this section, we explain how to compute the gradient of  $\log f(x)$ .
Lemma C.10. If the following condition holds
- Suppose that function  $f$  is defined in Definition C.4.
$$
\frac {\mathrm {d} \log f (x)}{\mathrm {d} x _ {i}} = - \langle f (x), A _ {*, i} \rangle \cdot \mathbf {1} _ {n} + A _ {*, i}
$$
$$
\langle \frac {\mathrm {d} \log f (x)}{\mathrm {d} x _ {i}}, b \rangle = \left\langle A _ {*, i}, b \right\rangle - \left\langle f (x), A _ {*, i} \right\rangle \cdot \left\langle b, \mathbf {1} _ {n} \right\rangle
$$
$$
\frac {\mathrm {d}}{\mathrm {d} x _ {i}} L _ {\mathrm {c e n t}} (x) = \langle f (x), A _ {*, i} \rangle \cdot \langle b, \mathbf {1} _ {n} \rangle - \langle A _ {*, i}, b \rangle
$$
Proof. Proof of Part 1.
For all index  $j \in [n]$ , we can compute the gradient with respect to  $x_{i}$
$$
\frac {\mathrm {d} \log f (x) _ {j}}{\mathrm {d} x _ {i}} = f (x) _ {j} ^ {- 1} \frac {\mathrm {d} f (x) _ {j}}{\mathrm {d} x _ {i}}
$$
Then we group the  $n$  coordinates, we get