TLDR: MemCom is a novel method that significantly improves the memory and computational efficiency of In-Context Learning (ICL) in Large Language Models (LLMs) by compressing many-shot prompts. Unlike previous methods, MemCom uses a powerful compressor and a unique layer-wise compression strategy, where compressed representations are provided to each transformer layer. This approach allows LLMs to maintain high accuracy even at high compression ratios (up to 8x), outperforming baselines that suffer significant performance drops. MemCom makes many-shot ICL more practical for deployment in resource-constrained environments.
Large Language Models (LLMs) have transformed how we approach various tasks, often learning new skills without explicit fine-tuning. This is largely thanks to a technique called In-Context Learning (ICL), where models are given several input-output examples, or “shots,” directly within the prompt. The more examples provided, the better the model performs. However, this benefit comes with a significant drawback: increasing the number of shots leads to much higher memory and computational costs during inference, creating a bottleneck for practical deployment, especially in resource-constrained environments like edge devices.
A new research paper, “Compressing Many-Shots in In-Context Learning,” introduces an innovative solution called MemCom to tackle this challenge. The core idea is to compress these lengthy “many-shot” prompts into a much smaller, more efficient summary without sacrificing performance. The paper highlights that existing prompt compression methods are often ineffective for this specific problem, primarily because they are designed for shorter prompts or only compress the final hidden states of the model. Interestingly, simply using fewer shots as a baseline proved to be a surprisingly strong competitor, underscoring the difficulty of effective compression.
The researchers identified two key insights for achieving effective compression: first, a more powerful compressor model with more trainable parameters is essential; and second, a fine-grained, layer-wise compression approach is necessary. Instead of compressing the entire prompt into a single summary at the end, MemCom compresses the representations at each transformer layer. This allows each layer of the main LLM to access its own tailored, compressed representation of the context, mimicking how the model would process the full, uncompressed input more closely.
MemCom’s architecture involves two stacks of LLMs working together as a compressor. A Source-LLM processes the original many-shot input tokens. A Memory-LLM then takes a fixed set of “memory tokens” and, at each layer, performs cross-attention over the corresponding layer representations from the Source-LLM. These resulting compressed memory representations are then passed to the Target-LLM, which is the model used for inference. Crucially, the Target-LLM remains frozen and attends only to these compressed memory tokens at each layer, drastically reducing memory usage and computational load. This compression process can be performed offline, making it ideal for scenarios where a powerful cloud system can pre-process prompts for less powerful edge devices.
The training of MemCom is conducted in two phases using large-scale pretraining datasets. Phase-1 focuses on training the randomly initialized cross-attention modules and memory tokens. Phase-2 then unfreezes and continues training the entire Source-LLM and Memory-LLM stacks. The Target-LLM remains frozen throughout, ensuring its broader capabilities are not interfered with.
Evaluations were conducted on various classification tasks with large label sets, using different model sizes (Gemma2-2B and Mistral-7B) and many-shot sequence lengths (3k-6k tokens), with compression ratios ranging from 3x to 8x. MemCom consistently outperformed strong baselines, including an enhanced version of a prior method called ICAE++. While baseline performance degraded sharply at higher compression ratios, often by over 20-30%, MemCom maintained high accuracy with minimal degradation, typically dropping by less than 10%. This robustness at higher compression ratios and longer sequence lengths highlights MemCom’s practical utility for real-world applications.
Also Read:
- Capsule Prompt Tuning: Efficiently Guiding Large Language Models
- Optimizing Vision-Language Models: A New Approach to Visual Token Pruning
While MemCom shows great promise, the authors acknowledge certain limitations. The method was primarily evaluated on classification tasks, and its impact on some generation tasks was less pronounced. The training cost for MemCom is substantial, requiring 80-160 billion tokens, which could increase with even longer sequences. Additionally, tasks requiring precise access to individual input tokens, such as low-resource language translation, might not benefit as much from this compression approach. Nevertheless, MemCom offers a significant step forward in making many-shot ICL more efficient and scalable. You can read the full paper here: Compressing Many-Shots in In-Context Learning.
