Efficient Many-Shot In-Context Learning with Dynamic Block-Sparse Attention

TL;DR

This work tackles the efficiency bottleneck of many-shot in-context learning by introducing Dynamic Block-Sparse Attention (DBSA), a training-free framework that combines block-sparse pre-encoding with dynamic retrieval and KV-cache reuse. By encoding demonstrations in blocks and selectively reusing precomputed caches for each query, DBSA achieves near-finettuning latency while preserving the high accuracy typical of many-shot ICL, as demonstrated on five datasets with Llama-2 and Llama-3.1 models and context lengths up to 90k tokens. The authors provide extensive ablations showing the importance of the anchor sink, block grouping, and ordering, and demonstrate robust performance gains across efficiency and accuracy metrics, suggesting practical deployment feasibility for scalable, retrieval-based ICL. Limitations include storage costs for KV caches and reliance on effective retrieval, with future directions focusing on adaptive retrieval strategies and cache management.

Abstract

Many-shot in-context learning has recently shown promise as an alternative to finetuning, with the major advantage that the same model can be served for multiple tasks. However, this shifts the computational burden from training-time to inference-time, making deployment of many-shot ICL challenging to justify in-practice. This cost is further increased if a custom demonstration set is retrieved for each inference example. We present Dynamic Block-Sparse Attention, a training-free framework for retrieval-based many-shot in-context learning. By combining carefully designed block-sparse attention and retrieval of cached groups of demonstrations, we achieve comparable per-example latency to finetuning while maintaining on average >95% of the best method's accuracy across strong ICL and finetuning baselines. We hope that this will further enable the deployment of many-shot ICL at scale.

Figures (6)

• Figure 1: DBSA maintains high accuracy while achieving the best overall efficiency when compared to many-shot ICL baselines and finetuning, even under high request volumes.
• Figure 2: DBSA pre-encodes the many-shot demonstration pool with block-sparse attention, and dynamically selects relevant KV chunks during inference for each test query.
• Figure 3: The total inference latency for each method, including setup and per-request inference latency. We use inference batch size 1.
• Figure 4: Storage cost for Cached many-shot ICL, which encompass Fixed ICL and DBSA
• Figure 5: Visualization of block-sparse attention mechanisms for many-shot ICL in section \ref{['sec:sparse attention']}.