Many-shot learning made faster and smaller with smarter memory use
MILO: Efficient Many-shot In-Context Learning with Block-wise Low-rank Compression
Computation and Language
Summary
Large language models can learn from many examples shown at once, but this uses a lot of computer memory, slowing things down. The authors proposed a way to compress the memory these models use by grouping examples and shrinking their stored data without losing important information. Their method adapts how much each group is compressed based on how much information it holds. Tests showed this approach halves the memory needs and nearly doubles speed with little drop in accuracy.
What this means in practice
- •For cloud ai engineers: Run large language models in many-shot mode using less memory and faster response by compressing key-value cache data efficiently.
- •For mobile app developers: Deploy large language models with many-shot learning on devices with limited memory by compressing the model’s context memories.
Authors
Youpeng Zhao, Tian Tan, Liqian Peng, Jun Wang, Alec Go
Abstract
Many-shot in-context learning (ICL) enables large language models (LLMs) to adapt to complex tasks by conditioning on thousands of demonstration examples, but this paradigm shifts the inference efficiency bottleneck to the key-value (KV) cache memory. Due to the linear scaling behavior of the KV cache, storing these intermediate tensors has become a paramount challenge for both online serving and on-device deployment. To address this issue, we propose a novel compression framework, termed MILO, that exploits the low-rank redundancy inherent in many-shot contexts. Specifically, MILO features a block-wise low-rank compression strategy that compresses the KV cache at the block granularity, where each block contains multiple many-shot examples. Furthermore, to handle the heterogeneous context density across different blocks, MILO dynamically allocates rank budgets based on the information entropy, preserving the fidelity of critical blocks while aggressively compressing redundant ones. Experimental results on Qwen2.5 models demonstrate that our method achieves up to 50% reduction in KV cache memory and 1.8x throughput improvement, with negligible performance degradation on classification and reasoning benchmarks, significantly outperforming prior baselines.