mlsanityData to 14 Oct 2026

Long-Context Demonstration Selection Using State Space Models

arXiv 2609.17888 ↗First posted 15 Sep 2026cs.LG cs.CLcited by 0

Not on the map. Why?

Abstract (arXiv)

We study the problem of demonstration selection, which involves selecting a subset of examples for prepending to a query to a language model. This problem is closely related to in-context learning and language model inference. Since the inference cost of a transformer model scales quadratically with sequence length, the selection problem becomes especially challenging in a long-context scenario. In this paper, we tackle this problem by building on state space models (SSMs), which require only linear inference time given the input. Our approach involves two algorithms. The first learns a small set of SSMs through distillation of a (trained) transformer model. We partition all the layers into consecutive groups. Then for each group, we estimate a separate state space model to replicate the input-output behavior within the adjacent layers. Second, we map the distilled model outputs to a small set of tokens, and apply these embeddings for demonstration selection in downstream applications. We perform extensive experiments in both synthetic and real-world datasets to validate our approach. We demonstrate that the distilled SSMs only incur an approximation error of less than 0.7% relative to the true output. In downstream evaluation, we show that on several text classification and reasoning tasks, our approach reduces FLOPs by 14.2× and improves accuracy by 6.48% relative to baseline demonstration selection methods.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 24 of its 31 references that are in the map. First posted on arXiv 15 Sep 2026.

Builds on

24 of 31 references.

References · 24

  1. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
  2. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  3. 2023Efficient Streaming Language Models with Attention Sinkscited by 2,186
  4. 2022Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?cited by 1,434
  5. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
  6. 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 605
  7. 2024Parallelizing Linear Transformers with the Delta Rule over Sequence Lengthcited by 396
  8. 2023Learning to Compress Prompts with Gist Tokenscited by 390
  9. 2023Adapting Language Models to Compress Contextscited by 356
  10. 2024DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Headscited by 288
  11. 2024Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modelingcited by 174
  12. 2025PaTH Attention: Position Encoding via Accumulating Householder Transformationscited by 34
12 more
  1. 2020Precise High-Dimensional Asymptotics for Quantifying Heterogeneous Transferscited by 27
  2. 2025From Few to Many: Self-Improving Many-Shot Reasoners Through Iterative Optimization and Generationcited by 21
  3. 2024Scalable Multitask Learning Using Gradient-based Estimation of Task Affinitycited by 20
  4. 2023Boosting Multitask Learning on Graphs through Higher-Order Task Affinitiescited by 15
  5. 2024Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approachcited by 13
  6. 2024Data-Driven Switchback Experiments: Theoretical Tradeoffs and Empirical Bayes Designscited by 12
  7. 2025Efficient Many-Shot In-Context Learning with Dynamic Block-Sparse Attentioncited by 12
  8. 2025Efficient Ensemble for Fine-tuning Language Models on Multiple Datasetscited by 10
  9. 2025Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimationcited by 8
  10. 2025Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimationcited by 8
  11. 2026Efficient Estimation of Kernel Surrogate Models for Task Attributioncited by 4
  12. 2024Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyondcited by 3

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Efficient and Effective In-context Demonstration Selection with Coreset
  2. 2026You Only Edit Once: Incentivizing In-Context Capability of LLMs via Local Demonstration Refinement
  3. 2024Large Language Models Know What Makes Exemplary Contexts
  4. 2025Leveraging In-Context Learning for Language Model Agents
  5. 2025Refract ICL: Rethinking Example Selection in the Era of Million-Token Models
  6. 2023In-Context Learning with Iterative Demonstration Selection
  7. 2024Demonstration Notebook: Finding the Most Suited In-Context Learning Example from Interactions
  8. 2024In-Context Learning with Reinforcement Learning for Incomplete Utterance Rewriting