Unifying Demonstration Selection and Compression for In-Context Learning
Authors: Jun Gao, Qian Qiao, Ziqiang Cao, Zili Wang, Wenjie Li


Reviewer paper summaries: The paper considers an in-context learning paradigm where the in-context demonstrations are selected and compressed into virtual tokens. Instead of having separate models for demonstration ranking, compression, and generation, the paper proposes to train a single model for all three tasks. The base generative LLM is kept fixed, and adaptation layers are trained for ranking and compression.

The compression allows the more demonstrations to fit in the context window, which is crucial for tasks like summarization. The experiments compare 5 types of methods:

1. S-BERT ranker + No compressor (normal in-context learning) + Vicuna generator
2. S-BERT ranker + a compressor that produces real tokens by dropping uninformative tokens + Vicuna generator. This is the LLMLingua baseline.
3. S-BERT ranker + a compressor that produces virtual tokens (specific to the generator) + its corresponding generator. This includes the AutoCompressors and ICAE baselines.
4. S-BERT ranker + proposed compressor + Vicuna generator
5. proposed ranker + proposed compressor + Vicuna generator
Summary of strengths: 1. **The method seems to work well compared to the baselines.**
    * Normal in-context learning (Method 1 above) can only fit 1 demonstration to the context window in some cases.
    * With the same number of shots, the proposed compressor (Method 4) also outperforms normal in-context learning. This means the compressed virtual tokens are actually better at guiding the target LLM than the actual tokens. The proposed compression method also outperforms the baseline compressors when the same ranker is used.
    * Replacing S-BERT with the proposed ranker also improves the performance further, though this is less surprising (a fine-tuned retriever is probably better than a pretrained retriever).
Summary of weaknesses: 1. (minor) **The motivations to train a single model are questionable.**
    * From Line 77, the paper claims that (1) having separate ranker and compressor lead to out-of-distribution problems between these extra models and the main LLM, and (2) the extra models incur additional memory costs.
    * For (1), the cited work ([Jiang et al., 2023](https://aclanthology.org/2023.emnlp-main.825/)) does mention that the compressor that is too different from the main LLM gives worse results. However, one could simply choose a model from the same family for training the compressor. Furthermore, Jiang23 compresses the prompt by removing tokens from the prompt, so the compressor needs to estimate how the main LLM would be affected by the removed tokens. In contrast, the compressor in this paper produces virtual token sequences (soft prompts), which are no longer normal tokens in the target LLM, so it might be less important to have matching models.
    * For (2), depending on the size of the extra models, the memory cost could be negligible. Admittedly, this motivation is still valid when there is only enough memory for one model, such as in an IoT device. There is also a tradeoff between memory and throughput, though this depends on the model architectures.

2. (major) **The novelty is limited.** As mentioned in the paper, the ideas of ranking and compressing demonstrations for in-context learning, as well as the methods for training such models, are based on previous works. The only contribution is the usage of a single model with trainable adaptation layers.

3. (major) **The baselines may not be exactly comparable to the proposed method.**
    * Method 2 above replaces the generative LLM with Vicuna. This is from the work mentioned in Weakness 1 above. So while the fine-tuned Llama 2 could generalize to a different target LLM, it might be better to use an in-distribution base model to prevent the out-of-distribution problem.
    * Method 3 above seems to use the corresponding target LLM for each compressor for generation. A more fair comparison would be to use Vicuna as the target LLM and re-train the compressor.
Comments, suggestions and typos: The writing is sometimes unclear and contains a lot of grammatical errors. Some errors up to Section 2 are given below.
- Line 1: “ICL facilitates LLMs exhibiting spectacular emergent capabilities” → “ICL enables LLMs to exhibit spectacular emergent capabilities.”
- Line 18: Remove “respectively”
- Line 55: “length disaster” sounds strange. Maybe rewrite it as something like “introducing multiple demonstrations can cause the input length to exceed the context window.”
- Line 61: “Except for” → “Apart from”
- Line 62: \citet
- Line 62: “of selected demonstrations” → “of the selected demonstrations”
- Figure 1 caption: Add something before the “*” in the last sentence.
- Line 69: “length explode” → “length explosion”
- Line 73: Remove “Similarly” since it’s not similar. Also use \citet.
- Line 101: “behind demonstrations” → “behind the demonstrations”
- Line 103: “on top of compressions slots” → “on top of the compression slots”
- Line 106: Not sure what “due to the transformer blocks” means.
- Line 108: “briefly explain” → “prompt”?
- Line 109: The linguistic acceptability task could be motivated more.
- Line 114: Remove the whole “along with inference inputs, fed … normal ICL”
- Line 118: “while” → “and”?
- Line 123: “Demonstrations Bank” → “Demonstration Bank”.
- Line 128: “equips” → “contains”
- Figure 2 contains multiple typos such as “Seleciton” and “Facnce”. Also the labels seem to be swapped (from CoLA, the first sentence is Acceptable while the second is not).
- Line 131: “out-domain” → “out-of-domain”
- Line 163 and Line 170: \citet
- Line 174: “practical” → “practice”

**Post-rebuttal comments**

The response addressed half of Weakness 1 and made a promise for the other half. With the remaining weaknesses and unclear writing, I want to keep my current score.
Overall assessment scores: 2.5
Metareview score: 2
Metareview: This paper addresses the problem of demonstration selection and compression in in-context learning by proposing a unified framework called UniICL. UniICL integrates demonstration selection, compression, and response generation into a single process, compressing demonstrations and text inputs into virtual tokens to prompt a language model. The method involves training adaptation layers for ranking and compression while keeping the base generative LLM fixed. Experiments on various classification and generation tasks across multiple datasets show that UniICL outperforms several baselines and generalizes well to new data and tasks.
Summary of reasons to publish: The method seems to work well compared to the baselines.

The method has a small number of trainable parameters.