Unifying Demonstration Selection and Compression for In-Context Learning

2024-05-27

arXiv.org

AI Summary

Existing in-context learning (ICL) methods suffer from two key bottlenecks: excessive prompt length—leading to high hardware overhead—and shallow, task-irrelevant demonstrations that impair model generalization. To address these, we propose UniICL, a unified framework that jointly optimizes demonstration selection, semantic compression, and response generation within a single frozen large language model (LLM). For the first time, UniICL integrates this three-stage co-optimization into a lightweight architecture with only 17M trainable parameters. Key innovations include virtual token projection, latent-space semantic similarity matching, and parameter-efficient fine-tuning of the frozen LLM. Extensive multi-task, multi-domain experiments demonstrate a 12× demonstration compression ratio. On IMDb, UniICL achieves significant accuracy gains under 64-shot ICL while requiring only 24GB CUDA memory.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Generation

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

In-context learning (ICL) facilitates large language models (LLMs) exhibiting spectacular emergent capabilities in various scenarios. Unfortunately, introducing demonstrations easily makes the prompt length explode, bringing a significant burden to hardware. In addition, random demonstrations usually achieve limited improvements in ICL, necessitating demonstration selection among accessible candidates. Previous studies introduce extra modules to perform demonstration compression or selection independently. In this paper, we propose an ICL framework UniICL, which Unifies demonstration selection and compression, and final response generation via a single frozen LLM. Specifically, UniICL first projects actual demonstrations and inference text inputs into short virtual tokens, respectively. Then, virtual tokens are applied to select suitable demonstrations by measuring semantic similarity within latent space among candidate demonstrations and inference input. Finally, inference text inputs together with selected virtual demonstrations are fed into the same frozen LLM for response generation. Notably, UniICL is a parameter-efficient framework that only contains 17M trainable parameters originating from the projection layer. We conduct experiments and analysis over in- and out-domain datasets of both generative and understanding tasks, encompassing ICL scenarios with plentiful and limited demonstration candidates. Results show that UniICL effectively unifies $12 imes$ compression, demonstration selection, and response generation, efficiently scaling up the baseline from 4-shot to 64-shot ICL in IMDb with 24 GB CUDA allocation

Problem

• Excessive context length growth in ICL burdens hardware
• Shallow-relevant examples hinder useful contextual information capture
• Need unified framework for compression, selection, and generation

Innovation

• Unified framework for compression, selection, generation
• Demonstration Bank caches compressed results
• Tailored compression strategy boosts efficiency