PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tsz Ting Chung, Leyang Cui, Lemao Liu, Xinting Huang, Shuming Shi, Dit-Yan Yeung

arXiv:2410.11786v2cs.CLcs.AIcs.LG

TL;DR

In-context learning lengthens prompts and raises computational and financial costs, while existing compression methods can be model-specific or rely on costly external training data. Selection-p trains with a self-supervised objective to predict tokens from compressed context and uses token-level probabilities to select what to preserve. Our model surpasses existing state-of-the-art compression models; at a 10x compression rate, it shows a 0.8-point drop across traditional classification tasks with a 5.3x speedup.

  • Problem

    In-context learning lengthens prompts and raises computational and financial costs, while existing compression methods can be model-specific or rely on costly external training data.

    Source: §1 · p. 1

  • Method

    Selection-p trains with a self-supervised objective to predict tokens from compressed context and uses token-level probabilities to select what to preserve.

    Source: §3 · p. 2, §1 · p. 2, §3.2 · p. 3

  • Results

    Our model surpasses existing state-of-the-art compression models; at a 10x compression rate, it shows a 0.8-point drop across traditional classification tasks with a 5.3x speedup.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Long-context experiments present the possibility of addressing the in-context learning issue of recent long-context models.

    Source: §6 · p. 9

    Due to cost, the authors did not analyze models beyond ChatGPT and Gemini, and LLaMA-2-7B does not achieve better latency than LLMLingua-2 or AutoCompressor.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large Language Models (LLMs) have demonstrated impressive capabilities in a wide range of natural language processing tasks when leveraging in-context learning. To mitigate the additional computational and financial costs associated with in-context learning, several prompt compression methods have been proposed to compress the in-context learning prompts. Despite their success, these methods face challenges with transferability due to model-specific compression, or rely on external training data, such as GPT-4. In this paper, we investigate the ability of LLMs to develop a unified compression method that discretizes uninformative tokens, utilizing a self-supervised pre-training technique. By introducing a small number of parameters during the continual pre-training, the proposed Selection-p produces a probability for each input token, indicating whether to preserve or discard it. Experiments show Selection-p achieves state-of-the-art performance across numerous classification tasks, achieving compression rates of up to 10 times while experiencing only a marginal 0.8% decrease in performance. Moreover, it exhibits superior transferability to different models compared to prior work. Additionally, we further analyze how Selection-p helps maintain performance on in-context learning with long contexts.

1 Introduction

In-context learning raises prompt costs, while existing compression methods face transferability challenges or rely on costly external training data. Selection-p introduces self-supervised token selection and is reported to compress prompts with limited performance loss while supporting transfer to other models and analysis in long-context settings.

  • In-context learning lengthens prompts and raises computational and financial costs, while prior compression can depend on GPT-4 training data or model-specific optimization that may hinder transfer.Continuous compression methods are also described as requiring fixed compression lengths and separate compressors for different models.Source: §1 · pp. 1–2
  • Selection-p uses a self-supervised pretraining signal to predict the next token from compressed context and produces token-level preserve-or-discard probabilities.Its small number of additional parameters supports applying detokenized compressed tokens to downstream LLMs with single-turn decoding and without costly external resources.Source: §1 · p. 2
  • At 10x compression, Selection-p incurs only a 0.8% performance drop across nine traditional classification tasks and surpasses existing compression models.The introduction also reports a 5.3x inference speedup under this setting.Source: §1 · p. 2
  • Selection-p is reported to transfer better than prior work under hard compression for both open-source and closed-source models.Source: §1 · p. 2
  • The paper analyzes Selection-p in long-context settings as a potential way to address performance decline in long-context in-context learning.Source: §1 · p. 2
  • The introduction positions the work through comparisons with in-domain prior work and state-of-the-art compression models.Source: §1 · p. 2

2 Related Work

Prior hard-compression work includes token pruning and merging, while soft-compression methods use continuous tokens. The paper contrasts Selection-p with both strands of work.

2.1 Hard Compression

  • Earlier token-pruning and token-merging methods were designed primarily for smaller models such as BERT; entropy-based pruning and LLMLingua extend hard compression to prompts and demonstrations.The passage notes that downstream efficacy with compressed in-context demonstrations remains limited for the no-training methods it describes.Source: §2.1 · p. 2

2.2 Soft Compression

  • Soft-token methods extend prompt compression to long contexts, but are described as using fixed compression rates and requiring retraining for different LLM versions.The paper presents its approach as more flexible and transferable than these methods while surpassing existing compression models.Source: §2.2 · p. 2

3 Methodology

The methodology trains a token-selection model to preserve language-modeling ability after context compression, then uses its token scores to pass a compressed context to a target LLM. Training uses a discretized selection mask and a CLM objective, with partial LoRA training intended to improve transferability.

3.2 Selection-p

  • The self-supervised training objective encourages the same token prediction before and after discarding less informative tokens, avoiding external training resources.The model discretizes scores into a binary top-k mask and applies a CLM loss to the masked context.Source: §3 · p. 2, §3.2 · p. 3

3.1 Preliminary

  • The preliminary setup describes causal language modeling as next-token prediction from prior context and notes that prior token selectors use the targeted LLM’s token distribution and a threshold.Source: §3.1 · p. 3

3.2 Selection-p

  • Selection-p separates compression from inference: a selection model chooses less informative context tokens, which are then passed to the targeted inference model.Source: §3.2 · p. 3
  • A pretrained LM’s last-layer hidden states feed a linear projection that assigns per-token informativeness scores; a single forward pass retains the top k% for inference.The sigmoid projection uses parameters W and b.Source: §3.2 · p. 3
  • To support transfer, the method freezes the pretrained LM and uses LoRA to train only part of the language model used in the masked-context loss.Figure 1 depicts an initial inference pass that creates the attention mask before LoRA training updates the additional linear layer.Source: Figure 1 · p. 3, §3.2 · p. 4

4 Experiment

The experiments assess Selection-p on traditional and long-context classification and test transfer to other LLMs. Selection-p leads the compared compression models on traditional tasks and outperforms tested long-context and transfer baselines.

4.1 Setting

  • Selection-p was evaluated on nine traditional classification tasks and BANKING77 long-context classification after LoRA finetuning on 100M RedPajama tokens in 1,024-token segments.For each task, demonstrations were sampled from the training set; compression was computed once for subsequent test inferences.Source: §4.1 · p. 4

4.2 Baselines

  • The comparisons included token-level LLMLingua, GPT-4-distilled LLMLingua-2, and RMT-based AutoCompressor, the sole soft-compression baseline.For the Table 2 comparison, the authors used LLMLingua’s token-level algorithm.Source: §4.2 · p. 5

4.3 Evaluation Result

  • Selection-p achieved the highest ICL performance among compression models, averaging 67.4% accuracy across all 10 tasks.No compression model exceeded full-shot overall; on WSC, hard compression models surpassed zero-shot and full-shot by approximately 20%, and Selection-p beat a one-tenth demonstration baseline with a comparable token count.Source: §4.3 · p. 5, Table 2 · p. 6
  • On BANKING77 long-context classification, Selection-p outperformed the compared long-context models and LLMLingua-2 at every token-size level.Its results also showed a growing trend as compressed demonstrations increased, indicating that it could learn from additional information after compression.Source: §4.3 · p. 5, Table 3 · p. 6

4.4 Transferability

  • On LLaMA-2-13B, Selection-p outperformed the other compression models, with only a small deviation from full-shot at 10x compression.Tasks that exceeded full-shot with Selection-p on LLaMA-2-7B showed similar patterns on LLaMA-2-13B.Source: §4.4 · p. 5
  • At 3x compression on BANKING77 using EM, Selection-p achieved the best performance among existing methods on ChatGPT while remaining below full-shot.The black-box evaluation used 750-token prompts; discrete compressed tokens sometimes triggered SAFETY errors.Source: §4.4 · pp. 5–6

5 Analysis

The analysis tests flexibility across demonstration lengths and models, measures latency, and examines token-level preservation and comparisons with LLMLingua. Selection-p surpasses LLMLingua with a Budget Controller, while LLMLingua’s achieved token size exceeds its target.

5.1 Flexibility

  • At about 250 tokens, Selection-p performs best among the compared compression models, and its average accuracy rises with more demonstrations.Other compression models did not improve with more demonstrations; AutoCompressor dropped 2% after an additional 500 tokens.Source: §5.1 · p. 6

5.2 Latency Analysis

  • Selection-p achieves a 5.3x speedup on 10x-compressed WSC demonstrations, with negligible compression time in the ICL setting.The latency analysis uses an A100-80G GPU.Source: §5.2 · p. 6, Table 7 · p. 8

5.3 Correlation with Attention and Perplexity

  • Selection-p’s token-preservation probability is only weakly correlated with mean attention and token-level perplexity, whose correlation is more significant.The analysis uses last-layer mean attention and Spearman’s rank correlation across traditional classification tasks.Source: §5.3 · p. 6, Figure 2 · p. 7

5.4 Tokens Level Part-of-Speech Analysis

  • Selection-p more often preserves PRPs and punctuation, while noun phrases are preserved more than verbs; formatting tokens such as colons are less likely to be preserved.The authors suggest PRP preservation may relate to WSC pronoun resolution and that punctuation helps retain sentence boundaries.Source: §5.4 · p. 7

5.5 On Fair Comparison with LLMLingua

  • With a Budget Controller, Selection-p surpasses LLMLingua on WSC, while both show a performance decrease.The comparison considers Selection-p compression rates of 10x and 38x.Source: §5.5 · p. 8, Table 8 · p. 8
  • LLMLingua targets 75 compressed tokens but averages about 192.1, exceeding the desired size by more than 1.5 times.The reported disparity spans the classification tasks.Source: §5.5 · p. 8

6 Conclusion

The paper presents a self-supervised approach to context compression and evaluates it across classification and long-context settings. Selection-p reports a 0.8-point performance drop at 10x compression with a 5.3x speedup, alongside transfer to open- and black-box models.

6 Conclusion

  • At 10x compression, Selection-p shows a 0.8-point drop in performance across traditional classification tasks with a 5.3x speedup.The evaluation also covers few-shot and long-context settings.Source: §6 · p. 9
  • Selection-p surpasses existing state-of-the-art compression models in evaluations including LLaMA-2-13B, ChatGPT, and Gemini.The paper describes transfer to both open-source and black-box models.Source: §6 · p. 9
  • Long-context experiments present the possibility of addressing in-context learning performance decline in recent long-context models.Source: §6 · p. 9

Limitations

The authors identify limits in model coverage and latency, while noting that compression time is minimal in the ICL setting.

Limitations

  • Cost constraints limited further model analysis to ChatGPT and Gemini beyond the evaluated LLaMA-2 models.Source: Limitations · p. 9
  • The LLaMA-2-7B-based Selection-p does not achieve better latency than LLMLingua-2 or AutoCompressor.Source: Limitations · p. 9
  • AutoCompressor’s soft compression limits its applicability to other LLMs, despite its better latency.Source: Limitations · p. 9

A Examples Illustration of Compression Results

The section presents compressed-prompt examples for two traditional classification tasks and one additional task.

A Examples Illustration of Compression Results

B Prompt of Evaluation on BANKING77

The BANKING77 evaluation prompt lists all 77 labels before a set of demonstrations, then uses a context-and-answer template.

B Prompt of Evaluation on BANKING77

C Training Details

Selection-p compression uses LLaMA-2-7B and is trained on 100M tokens from RedPajama, taking roughly 50 hours on one A100 GPU.

C Training Details

D Detailed Latency Analysis

For WSC, token-selection compression time is presented separately, while its contribution to end-to-end inference is described as negligible relative to inference. The section also illustrates compressed demonstrations for several tasks.

D Detailed Latency Analysis

  • Table 9 presents compression time per demonstration set for token-based compression models on WSC.Source: D Detailed Latency Analysis · p. 12, Table 9 · p. 12
  • The same demonstration set is reused under ICL, so its compression is computed once for subsequent inferences.This contributes to short compression time in the complete end-to-end process.Source: D Detailed Latency Analysis · p. 12
  • Compression time is approximately 0.01 of inference time, so token-selection models have roughly the same end-to-end inference time.The WSC end-to-end timing results are shown in Table 10, ordered by inference time for clarity.Source: D Detailed Latency Analysis · p. 12, Table 10 · p. 12
  • The examples include original Subj, WSC, and BANKING77 demonstrations, including BANKING77 context-answer pairs.The BANKING77 examples include cash_withdrawal_charge and Refund_not_showing_up.Source: D Detailed Latency Analysis · pp. 13–14
  • Figures 4 and 5 illustrate Selection-p compression at 10x for Subj, WSC, and BANKING77 demonstrations totaling about 750 tokens.The illustrations use 19 Subj demonstrations, 16 WSC demonstrations, and 27 BANKING77 demonstrations.Source: Figure 4 · p. 13, Figure 5 · p. 14

Related papers

Loading 2410.11786v2…