Selection-$p$: Self-Supervised Task-Agnostic Prompt Compression for Faithfulness and Transferability

Given a context $[x_{1},x_{2},\dots,x_{n-1}]$, the objective of the language model is to predict the next token $x_{n}$, formed as $P(x_{n}|x_{1},x_{2},\dots,x_{n-1})$. In case of training with the causal language modeling (CLM) loss, we will have,

LLMLingua Jiang et al. (2023) and its variant Li et al. (2023); Pan et al. (2024) select tokens according to the computed distribution of the targeted LLM. Specifically, suppose $x_{i}$ is a token in the prompt, if the probability of a token $x_{i}$ is less than a threshold, then the token is selected to be compressed.

By using the selection model $\mathbf{p}$, it is straightforward to compress the context for inference: we directly prune the corresponding tokens according to our desired compression rate in retaining the top $k\%$ of tokens in the context. To ensure the efficiency in performing compression, a single forward pass on our model creates the token preservation probability for all input tokens for simplicity.

Our training criterion for the selection model aims to preserve the language modeling ability of the LLM while also learning to discard tokens effectively. To this end, we employ the self-supervised approach to optimize the selection model and therefore we do not need external resources to train the selection model compared with  Pan et al. (2024).

One of our goals is to achieve a transferable compression method. Therefore, the parameters that achieve the best loss on the language models in Eq. 1 and Eq. 2 in training may not be transferred to the targeted language model in inference since the language models in training and inference can be different. As a result, to ensure the better transferability of the optimized selection model, we freeze the pre-trained language model in Eq. 1 and employ LoRA Hu et al. (2021) to train a partial parameter in the language model in Eq. 2. ¹¹1In our preliminary experiments, we tried the Gumbel-softmax trick to optimize the loss in Eq. 2 but we did not observe gains over the direct optimization implemented in our paper. In summary, the CLM-based training loss is illustrated in Figure 1.

Following Chevalier et al. (2023), we evaluate and compare different compression models on nine classification tasks, including six tasks from SuperGlue Wang et al. (2019). The predictions by LLMs are determined by iterating through all possible answer options for the instance and selecting the option with the minimum negative log-likelihood. The in-context demonstrations have been carefully selected to approximate a size of 750 tokens, and the complete demonstration is employed. This is referred to as the “full-shot”. Since the average token length for a single demonstration varies for different tasks (e.g., a single demonstration in RTE averages about 75 tokens, resulting in a 10-shot setup under the full-shot setting), the exact number of shots differs depending on the task.

To ensure a fair comparison among different compression models, a compression rate of 0.1 is used. For each task, four sets of demonstrations are selected, and the average result across these four trials is presented in Table 2. Additionally, the average accuracy of the nine tasks is presented for a clearer performance comparison.

To achieve a compression rate of 0.1, a simpler approach is to directly retain one-tenth of the full-shot demonstration (e.g., using 1-shot instead of 10-shot for the RTE task) instead of performing selection at the token level. We also include this method as a baseline to further demonstrate the effectiveness of our approach.

Recent research by Li et al. (2024) shows the failure of in-context learning tasks when applied to long-context scenarios. To investigate whether compression models can serve as a viable solution in long-context settings, we compare Selection-$p$ with long-context models, including LLaMA-2-7B-LongLora Chen et al. (2023) and Long-LLaMA-code-7B Tworkowski et al. (2023) on the BANKING77 dataset Casanueva et al. (2020). The dataset contains 77 classes where traversing all the instances with unique labels requires approximately two thousand tokens. Evaluation is conducted at 2K, 4K, and 7K token levels, and we adopt a compression rate of 10x for Selection-$p$ and LLMLingua-2 among all levels. Since the long in-context demonstration is used, chunking is performed for every 2,048 tokens in Selection-$p$. The compressed results are concatenated together with a space token between each pair of chunks. We again follow the evaluation setting by Chevalier et al. (2023) on the models’ prediction, with the result presented in Table 3.

None of the compression models can achieve superior performance compared to the full-shot demonstration setting, which is in line with our expectations given the information loss during compression. However, certain tasks show a notable improvement when compared to both the zero-shot and full-shot approaches, e.g., all the hard compression models surpass zero-shot and full-shot by approximately 20% in the WSC task. Among all compression models, Selection-$p$ demonstrates the highest performance in conducting ICL, with an average accuracy of 67.4% across all 10 tasks as presented in Table 2. Examples of in-context demonstration before and after compression are shown in Appendix A. To demonstrate the effectiveness of our model, we have included one-tenth of the original demonstration set as the baseline. Our model significantly outperforms the baseline with a comparable number of tokens, therefore highlighting the effect of performing compression at the token level.

Our model outperforms LLaMA-2-7B-LongLora, Long-LLaMA-code-7B and LLMLingua-2 at all token size levels, and achieves similar results to  Li et al. (2024)’s findings on the long-context models. In addition, a growing trend with variations is observed with increasing compressed demonstrations, indicating that our model can successfully learn from additional information after compression. Examples of in-context demonstration before and after compression are presented in Appendix A.

We follow the same setting of evaluation across different classification tasks in Section 4.3. To assess the transferability of the compression models, we compress demonstrations with Selection-$p$ and input the compressed demonstration tokens into LLaMA-2-13B. In comparing different compression models, since retraining is required for soft compression methods, no results can be obtained for AutoCompressor Chevalier et al. (2023). In the case of LLMLingua, token-level perplexity is calculated with LLaMA-2-13B instead of LLaMaA-2-7B in this experiment.

Our approach outperforms all other compression models as shown in Table 4. In addition, a small deviation is observed between the 10x compression rate and the full shot setting, demonstrating the great transferability of our models. Notably, with Selection-$p$, the tasks that outperform the full-shot setting in LLaMA-2-7b also exhibit similar patterns in LLaMA-2-13B.

Taking cost into consideration, we select ChatGPT OpenAI (2023) and Gemini Team (2023) for evaluation to examine its transferability to LLMs. Traditional classification tasks often have a simple nature and the potential issue of data contamination, leading to high accuracy and causing an insignificant evaluation. Therefore, we use BANKING77 Casanueva et al. (2020) for evaluation. Following a similar setup as described in Section 4.3, we adopt a token size of 750 for examination. However, in case the compression rate is too high, ChatGPT and Gemini are more likely to deviate from the instructions and provide task-irrelevant responses. Therefore, we adopt a compression rate of 3x and use the EM metric for this experiment given their black-box nature. Note that there may be variation in the result of Gemini since the discrete compressed tokens sometimes trigger the SAFETY error. The prompt used is presented in Appendix B.

Though the performance of our model still deviates from the full-shot setting, it achieved the best performance compared to the existing works as presented in Table 4, demonstrating fair transferability even in LLMs like ChatGPT. Surprisingly, though LLMLingua-2 is distilled from GPT-4, it exhibits poor generalization compared to other compression models.

The result in long context classification tasks in Section 4.3 shows the effectiveness of chunk-wise compression in long context. We further analyze if compression models work well in normal few-shot settings in classification tasks. In this experiment, the in-context demonstrations are selected with an approximate size of 250 tokens. The comparison to the result with the token size of 750 in Section 4.3 is presented in Table 6.

Selection-$p$ shows the best performance under the constraint of 250 tokens when compared to other compression models. Additionally, it also follows the full-shot (i.e. 750 tokens level) trend, the average performance across all classification tasks increases along with the number of provided demonstrations. On the contrary, other compression models didn’t achieve an improvement in accuracy with more demonstrations. For instance, there is a drop of 2% recorded with an additional 500 tokens of information for AutoCompressor.