Compressing Many-Shots in In-Context Learning

TL;DR

This paper addresses the memory and compute costs of long-context many-shot In-Context Learning (ICL) by introducing MemCom, a layer-wise compression framework that maps a long prompt of length $t$ to $m$ memory tokens ($m<t$) across transformer layers. MemCom employs two compressor stacks (Source-LLM and Memory-LLM) to generate per-layer compressed representations that the frozen Target-LLM attends to during inference, reducing KV-cache storage and attention computation. A two-phase training procedure on large pretraining data yields a strong compressor, with MemCom achieving robust performance across compression ratios from $3\times$ to $8\times$ on five classification tasks with large label spaces, often outperforming ICAE and other baselines, especially at high compression. The approach enables practical edge deployment and cloud-edge hybrid setups by allowing offline compression of demonstrations and inference on compressed memory, though it incurs substantial compressor training costs and is best demonstrated on classification benchmarks where long-context benefits are most pronounced.

Abstract

Large Language Models (LLMs) have been shown to be able to learn different tasks without explicit finetuning when given many input-output examples / demonstrations through In-Context Learning (ICL). Increasing the number of examples, called ``shots'', improves downstream task performance but incurs higher memory and computational costs. In this work, we study an approach to improve the memory and computational efficiency of ICL inference by compressing the many-shot prompts. Given many shots comprising t tokens, our goal is to generate a m soft-token summary, where m < t. We first show that existing prompt compression methods are ineffective for many-shot compression, and simply using fewer shots as a baseline is surprisingly strong. To achieve effective compression, we find that: (a) a stronger compressor model with more trainable parameters is necessary, and (b) compressing many-shot representations at each transformer layer enables more fine-grained compression by providing each layer with its own compressed representation. Based on these insights, we propose MemCom, a layer-wise compression method. We systematically evaluate various compressor models and training approaches across different model sizes (2B and 7B), architectures (Gemma and Mistral), many-shot sequence lengths (3k-6k tokens), and compression ratios (3x to 8x). MemCom outperforms strong baselines across all compression ratios on multiple classification tasks with large label sets. Notably, while baseline performance degrades sharply at higher compression ratios, often by over 20-30%, MemCom maintains high accuracy with minimal degradation, typically dropping by less than 10%.

Figures (4)

• Figure 1: MemCom consists of two stacks of LLMs working as compressor. The Source-LLM processes the input tokens to be compressed. The Memory-LLM processed memory tokens, which cross-attends to the input representation in corresponding layer. Finally, during inference, Target-LLM (frozen ) processes test input and attends to the compressed representations at each layer. Training happens in two phases. In Phase-1, only the cross-attention module and the memory tokens are trained (represented by *X), while in Phase-2 the entire stack of Source-LLM and Memory-LLM is trained (*X*X), in addition to the memory tokens.
• Figure 2: Comparison of compression performance across different compression ratios. As the compression ratio increases, the vanilla baseline ("Baseline") exhibits significant performance degradation. ICAE++ performs better than the vanilla baseline but remains less robust compared to MemCom, which maintains high performance even at high compression ratios. The top row shows results for Mistral-7B, where sequences of $6k$ tokens are compressed. The bottom row shows results for Gemma2-2B, with $3k$-token sequences being compressed. MemCom refers to our method's performance post Phase-1 training, which introduces comparatively lightweight compressor of 1-head cross attention in every layer. While MemCom-P2 refers to the performance post Phase-2 training, which has two full stacks of LLMs working as the compressor. Section \ref{['sec:memcom_arch']} introduces the training methodology.
• Figure 3: (a) Architecture of ICAE - The compressor is a copy of the decoder LLM, with LoRA parameters applied only to the Query and Key matrices. We extend this setup in ICAE+ by including LoRA parameters to value and post matrices as well. And we eventually extend to the entire attention module being trained, referred as ICAE++. (b) Performance improves as we move from ICAE$\rightarrow$ICAE+$\mathbf{\rightarrow}$ICAE++$\rightarrow$MemCom. The comparison is shown through performance on Trec-Fine across training steps, on Mistral-7B while compressing 6k source tokens by $8\times$ into $768$ memory tokens.
• Figure 4: