Long-Context Demonstration Selection Using State Space Models

# Long-Context Demonstration Selection Using State Space Models
Source: [https://arxiv.org/html/2609.17888](https://arxiv.org/html/2609.17888)
###### Abstract

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 than0\.7%0\.7\\%relative to the true output\. In downstream evaluation, we show that on several text classification and reasoning tasks, our approach reduces FLOPs by14\.2×14\.2\\timesand improves accuracy by6\.48%6\.48\\%relative to baseline demonstration selection methods\.

## 1Introduction

Language models increasingly rely on answering to long\-context prompts at inference time\([Oncescu et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib7)\)\. One approach is to condition the model on demonstrations drawn from historical trajectories[Xiong \(2025\)](https://arxiv.org/html/2609.17888#bib.bib22)\. Specifically, given a query, the goal is to select a subset ofkkdemonstrations from a large candidate pool of sizenn\. This problem is known asdemonstration selection, which is closely related to in\-context learning[Garg et al\. \(2022\)](https://arxiv.org/html/2609.17888#bib.bib1);[Zhang et al\. \(2025b\)](https://arxiv.org/html/2609.17888#bib.bib2)\.

Figure 1:Overview of our approach: Given a set of candidate demonstrations and query\-target pairs from an unknown distribution, our approach learns a set of state space models through distillation from a pretrained model\. Then, we use the extracted embeddings from the state space models for demonstration selection\. We sample multiple subsets of demonstrations and evaluate the concatenated subset embeddings\. We aggregate the subset losses to an affinity score for every demonstration and select the top\-kkdemonstrations ranked by the scores\.The challenge in selecting demonstrations is that the effect of one demonstration often depends on which other demonstrations appear with it in the prompt\. Evaluating candidate subsets often requires repeatedly running the language model on many different prompts\([Zhang et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib2)\)\. For transformer models, the cost of processing long prompts grows quadratically with the total context lengthTT, whereTTscales with both the number and length of demonstrations\. This creates a computational bottleneck for long\-context demonstration selection, since exhaustively evaluating many candidate subsets is infeasible askkincreases\.

This challenge connects demonstration selection to a broader problem of efficient long\-context inference\. Existing approaches largely fall into three categories\. First, state space models \(SSMs\)\([Gu et al\., 2020](https://arxiv.org/html/2609.17888#bib.bib5);[Gu et al\., 2022](https://arxiv.org/html/2609.17888#bib.bib4);[Gu and Dao, 2024](https://arxiv.org/html/2609.17888#bib.bib9)\)and related hybrid architectures\([Ren et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib8);[Oncescu et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib7)\)achieve efficient sequence modeling, but typically require retraining the underlying language model\. Second, training\-free methods accelerate inference by restricting attention to local windows or selected tokens\([Xiao et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib6);[Xiao et al\., 2025a](https://arxiv.org/html/2609.17888#bib.bib12)\)\. Third, other related methods\([Mu et al\., 2023](https://arxiv.org/html/2609.17888#bib.bib15);[Chevalier et al\., 2023](https://arxiv.org/html/2609.17888#bib.bib16);[Ge et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib14)\)require processing the entire prompt\.

In this paper, we build on state space models to scale up demonstration selection to long\-contexts at inference time\. The key idea is to train a group of SSMs to map each candidate subset’s concatenated demonstrations into a small set of tokens\. These tokens are computed once per subset and reused across all queries evaluated with that subset\. This substantially reduces the cost of evaluating many candidate demonstration subsets\. Our algorithm consists of three main components\.

First, we learn a small number of state space models through distillation\. We partition the transformer layers into disjoint groups of consecutive layers and assign one SSM to each group\. Each SSM scans the embeddings of the long demonstration prompt into a compact hidden representation, which is then mapped into a small set of key\-value states\. The output key\-value states are then combined together as the tokens that transformers can process directly\. These tokens serve as a substitute for the original long prompt during inference\.

Second, we design a demonstration selection algorithm based on subset sampling[Li et al\. \(2023b\)](https://arxiv.org/html/2609.17888#bib.bib21);[Li et al\. \(2023a\)](https://arxiv.org/html/2609.17888#bib.bib28);[Li et al\. \(2024a\)](https://arxiv.org/html/2609.17888#bib.bib18);[Li et al\. \(2024b\)](https://arxiv.org/html/2609.17888#bib.bib19);[Zhang et al\. \(2025b\)](https://arxiv.org/html/2609.17888#bib.bib2)\. Crucially, this algorithm now runs on top of the \(distilled\) state space models\. We repeatedly sample candidate subsets and evaluate their predictive performance using the representation produced by our approach\. We then estimate the contribution of each demonstration based on how subsets containing that demonstration perform\.

Finally, we conduct empirical analyses of the proposed framework\. We find that our approach preserves the predictive distribution of the original model with relative error below0\.6%0\.6\\%\. Across classification and reasoning benchmarks, our approach reduces FLOPs by up to14\.2×14\.2\\timesand improves downstream accuracy by6\.48%6\.48\\%over baselines\.

In summary, we design an algorithm to scale up demonstration selection to long\-context inference\. First, we build SSMs that map long demonstration prompts into a small set of reusable tokens\. Second, we develop an affinity\-based subset selection algorithm that leverages this representation to evaluate many candidate demonstration subsets\. Third, we validate the proposed algorithms on synthetic, classification, and reasoning benchmarks\. The code to reproduce our experiments is available at[https://github\.com/VirtuosoResearch/Long\-context\-demonstration\-selection](https://github.com/VirtuosoResearch/Long-context-demonstration-selection)\.

## 2Preliminaries

We study demonstration selection in constructing a prompt\. LetDDdenote the set of candidate demonstrations of sizenn\. Let\(q,y\)∼𝒫\(q,y\)\\sim\\mathcal\{P\}denote a query\-target pair drawn from an unknown distribution\. We select a subsetD′⊆DD^\{\\prime\}\\subseteq Dcontainingkkdemonstrations with a fixed ordering for every query\. Given a modelfWf\_\{W\}and a loss functionℓ:𝒳×𝒴→\\ell:\\mathcal\{X\}\\times\\mathcal\{Y\}\\rightarrow\\realsupported on an input domain𝒳\\mathcal\{X\}and output domain𝒴\\mathcal\{Y\}, the*demonstration \(subset\) selection problem*is defined as the following minimization:

minD′⊆D:\|D′\|=k𝔼\(q,y\)∼𝒫\[ℓ\(fW\(D′,q\),y\)\]\.\\min\_\{\\begin\{subarray\}\{c\}D^\{\\prime\}\\subseteq D:\\,\\left\\lvert D^\{\\prime\}\\right\\rvert=k\\end\{subarray\}\}\\mathbb\{E\}\_\{\(q,y\)\\sim\\mathcal\{P\}\}\\left\[\\ell\\\!\\left\(f\_\{W\}\(D^\{\\prime\},q\),y\\right\)\\right\]\.
For the selected demonstrationsD′D^\{\\prime\}, we concatenated them as the input demonstration prompt\. LetX=\[x1,…,xT\]⊤∈ℝT×HX=\[x\_\{1\},\\dots,x\_\{T\}\]^\{\\top\}\\in\\mathbb\{R\}^\{T\\times H\}denote the token embeddings of the prompt, whereTTis the total number of tokens andHHis the embedding dimension\. SinceTTincreases with both the number and length of demonstrations, long\-context prompting can produce very large input sequences\. In a transformer model, the computational cost of processing the prompt grows quadratically withTTbecause self\-attention computes pairwise interactions between all tokens\. As a result, evaluating the inference outcome for a largeTTis computationally expensive\.

A natural approach for processing long prompts is to use state space models \(SSMs\), which summarize the input sequence through recurrent updates\. Given the prompt embeddingsXX, a discrete SSM updates its state at each token positionttas

ht=A¯​ht−1\+B¯​xt,h\_\{t\}=\\bar\{A\}h\_\{t\-1\}\+\\bar\{B\}x\_\{t\},\(1\)whereht∈ℝNh\_\{t\}\\in\\mathbb\{R\}^\{N\}is theNN\-dimensional state,A¯∈ℝN×N\\bar\{A\}\\in\\mathbb\{R\}^\{N\\times N\}is the state\-transition matrix, andB¯∈ℝN×H\\bar\{B\}\\in\\mathbb\{R\}^\{N\\times H\}maps the token embeddingxtx\_\{t\}into the state space\. The SSM processes each token once and has linear complexity in the sequence lengthTT[Gu et al\. \(2022\)](https://arxiv.org/html/2609.17888#bib.bib4)\.

Previous work has shown that the HiPPO\-LegS matrix provides a structured construction of the state\-transition matrix by maintaining an online polynomial projection of the input history\([Gu et al\., 2020](https://arxiv.org/html/2609.17888#bib.bib5);[Gu and Dao, 2024](https://arxiv.org/html/2609.17888#bib.bib9)\)\. The state approximates the coefficients of past inputs under a scaled Legendre basis\. Lower\-order coefficients capture the main structure of the sequence history, and higher\-order coefficients represent finer variations\. Thus, a finite\-dimensional HiPPO state provides a compact summary of a long input sequence\.

In the in\-context learning setting, we hypothesize that task\-relevant information in a large demonstration set can be preserved by low\-order Legendre basis components of the HiPPO matrix\. The SSM scans the complete input once\. Thus, it can replace a transformer model’s inference for the demonstrations and reduce the inference complexity fromO⁡\(T2\)\{O\}\(T^\{2\}\)toO⁡\(T\)\{O\}\(T\)\.

Motivated by these existing results, a natural question is whether one could adapt SSMs to tackle the demonstration subset selection problem\.

- •First, how can we design the SSM for the in\-context learning setting?
- •Second, how can we utilize the SSM for demonstration selection?

In the next section, we design algorithms to answer each of the above questions, respectively\.

## 3Our Approach

We build SSMs for efficient long\-context inference\. First, we present the architecture of the SSMs with theoretical and empirical evidence supporting our hypothesis\. We also introduce a distillation procedure for training the SSMs\. Then, we propose a demonstration selection method on the embeddings and use controlled synthetic experiments to evaluate its affinity estimates\.

### 3\.1Learning the State Space Models

We design SSMs with the HiPPO\-LegS matrix to map long demonstration prompts while preserving the dominant low\-order structure of the sequence\. In a transformer, the prompt is first processed to construct a layer\-wise key\-value \(KV\) cache, where each transformer layer stores key and value representations for all prompt tokens\. These cached representations are later used during decoding\.

Instead of computing the exact KV cache for the full prompt, we use SSMs to process it\. The SSM sequentially scans the prompt and summarizes the long sequence into a compact hidden representation\. We partition themmlayers of the transformer modelffintoggdisjoint groups of consecutive layers to yield a lower\-dimensional target for each group\. Each group is assigned an independent SSM that predicts the KV states only for its corresponding group of layers\.

Specifically, for each layer groupGiG\_\{i\}fori=1,…,gi=1,\\dots,g, the SSM processes the prompt embeddingsXXand produces a final hidden statehT\(i\)∈ℝNh\_\{T\}^\{\(i\)\}\\in\\mathbb\{R\}^\{N\}\. A multilayer perceptronfW\(i\)f\_\{W\}^\{\(i\)\}then projectshT\(i\)h\_\{T\}^\{\(i\)\}into a separate pair of key\-value states for every transformer layer in the group:

\{\(K~\(l\),V~\(l\)\)\}l∈Gi=fW\(i\)​\(hT\(i\)​\(X\)\)\.\\left\\\{\\bigl\(\\tilde\{K\}^\{\(l\)\},\\tilde\{V\}^\{\(l\)\}\\bigr\)\\right\\\}\_\{l\\in G\_\{i\}\}=f\_\{W\}^\{\(i\)\}\(h\_\{T\}^\{\(i\)\}\(X\)\)\.For each transformer layerll, the original KV states have shapeT×nkv×d\{T\\times n\_\{\\mathrm\{kv\}\}\\times d\}, while the projected states have shapenv×nkv×d\{n\_\{v\}\\times n\_\{\\mathrm\{kv\}\}\\times d\}, wherenkvn\_\{\\mathrm\{kv\}\}is the number of key\-value heads andddis the dimension of each head\. The KV states generated from the outputs of allggSSMs are aligned at the samenvn\_\{v\}token positions and assembled across layer groups to form the layer\-wise KV cache of thesenvn\_\{v\}tokens\. Thus, the projection preserves the key\-value head and head\-dimension axes while mappingTTtokens tonvn\_\{v\}tokens for a transformer to process directly\.

To analyze the proposed architecture, we first run the transformer on the complete input prompt using a standard full\-attention forward pass and record the resulting key\-value states at every layer\. We represent the exact transformer KV cache of each layer groupGiG\_\{i\}by a matrixC\(i\)C^\{\(i\)\}\. This matrix stacks the key and value states from all layers and attention heads in the group, and retains theTTtoken positions as its rows\.

###### Proposition 3\.1\.

Let\{ϕr\}r=1T\\\{\\phi\_\{r\}\\\}\_\{r=1\}^\{T\}be an orthonormal discrete Legendre basis along the token dimension, so thatϕr⊤​ϕr′=𝟙r=r′\\phi\_\{r\}^\{\\top\}\\phi\_\{r^\{\\prime\}\}=\\mathbbm\{1\}\_\{r=r^\{\\prime\}\}\. For eachGiG\_\{i\}, letC\(i\)∈ℝT×DiC^\{\(i\)\}\\in\\mathbb\{R\}^\{T\\times D\_\{i\}\}, whereDiD\_\{i\}is the number of stacked cache features, andcr\(i\):=\(C\(i\)\)⊤ϕr∈ℝDi,r=1,…,Tc\_\{r\}^\{\(i\)\}:=\\bigl\(C^\{\(i\)\}\\bigr\)^\{\\top\}\\phi\_\{r\}\\in\\mathbb\{R\}^\{D\_\{i\}\},r=1,\\ldots,T\. The orthogonal decomposition ofC\(i\)C^\{\(i\)\}is∑r=1Tϕr​\(cr\(i\)\)⊤\\sum\_\{r=1\}^\{T\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}\. ForN<TN<T, define the orthogonal projection matrix of rankNNontospan⁡\{ϕ1,…,ϕN\}\\operatorname\{span\}\\\{\\phi\_\{1\},\\ldots,\\phi\_\{N\}\\\}byΠN:=∑r=1Nϕr​ϕr⊤∈ℝT×T\.\\Pi\_\{N\}:=\\sum\_\{r=1\}^\{N\}\\phi\_\{r\}\\phi\_\{r\}^\{\\top\}\\in\\mathbb\{R\}^\{T\\times T\}\.It follows thatΠN​C\(i\)=∑r=1Nϕr​\(cr\(i\)\)⊤\\Pi\_\{N\}C^\{\(i\)\}=\\sum\_\{r=1\}^\{N\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}andC\(i\)−ΠN​C\(i\)=∑r=N\+1Tϕr​\(cr\(i\)\)⊤C^\{\(i\)\}\-\\Pi\_\{N\}C^\{\(i\)\}=\\sum\_\{r=N\+1\}^\{T\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}\. Therefore, for anys\>0s\>0,

‖C\(i\)−ΠN​C\(i\)‖F2≤∑r=N\+1Tr2​s\(N\+1\)2​s​‖cr\(i\)‖22\.\\displaystyle\\left\\\|C^\{\(i\)\}\-\\Pi\_\{N\}C^\{\(i\)\}\\right\\\|\_\{F\}^\{2\}\\leq\\sum\_\{r=N\+1\}^\{T\}\\frac\{r^\{2s\}\}\{\(N\+1\)^\{2s\}\}\\left\\\|c\_\{r\}^\{\(i\)\}\\right\\\|\_\{2\}^\{2\}\.

This shows that the error stems from Legendre components omitted by the firstNNmodes\. When the coefficients are concentrated at low frequencies, this tail is small\. Increasing the SSM state dimensionNNretains more components and further reduces the approximation error\. The proof is provided in Appendix[A](https://arxiv.org/html/2609.17888#A1)\.

Algorithm 1Learning SSMs via Key, Value, and Output DistillationInput: Embedding and query pairs\{X,q\}\\\{X,q\\\}, parameter initializationΘ=\{B¯\(i\),W\(i\)\}i=1g\\Theta=\\\{\\bar\{B\}^\{\(i\)\},W^\{\(i\)\}\\\}\_\{i=1\}^\{g\} Require: Groups\{Gi\}i=1g\\\{G\_\{i\}\\\}\_\{i=1\}^\{g\},ggSSMs each with a fixedA¯\\bar\{A\}and variable weight matrixB¯\(i\)\\bar\{B\}^\{\(i\)\}and MLPs\{fW\(i\)\}i=1g\\\{f\_\{W\}^\{\(i\)\}\\\}\_\{i=1\}^\{g\}, anmm\-layer transformer,nvn\_\{v\}, parametersλ1,λ2\\lambda\_\{1\},\\lambda\_\{2\}, and learning rateη\\eta Output: Trained parametersΘ^\\hat\{\\Theta\}

1:foreach

\(X,q\)\(X,q\)do

2:Extract

\{K\(l\),V\(l\)\}l=1m\\\{K^\{\(l\)\},V^\{\(l\)\}\\\}\_\{l=1\}^\{m\}from transformer

3:

T←T\\leftarrowLength of

XX
4:for

i∈\{1,…,g\}i\\in\\\{1,\\ldots,g\\\}do

5:

\{\(K~\(l\),V~\(l\)\)\}l∈Gi←fW\(i\)​\(hT\(i\)​\(X\)\)\\\{\(\\tilde\{K\}^\{\(l\)\},\\tilde\{V\}^\{\(l\)\}\)\\\}\_\{l\\in G\_\{i\}\}\\leftarrow f\_\{W\}^\{\(i\)\}\(h\_\{T\}^\{\(i\)\}\(X\)\)
6:endfor

7:for

i∈\{1,…,nv\}i\\in\\\{1,\.\.\.,n\_\{v\}\\\},

l∈\{1,…,m\}l\\in\\\{1,\\dots,m\\\}do

8:

κ1\(i,l\)←1−cos⁡\(K~i\(l\),Pool⁡\(K\(l\)\)i\)\\kappa\_\{1\}^\{\(i,l\)\}\\leftarrow 1\-\\cos\(\\tilde\{K\}^\{\(l\)\}\_\{i\},\\operatorname\{Pool\}\(K^\{\(l\)\}\)\_\{i\}\)
9:

κ2\(i,l\)←1−cos⁡\(V~i\(l\),Pool⁡\(V\(l\)\)i\)\\kappa\_\{2\}^\{\(i,l\)\}\\leftarrow 1\-\\cos\(\\tilde\{V\}^\{\(l\)\}\_\{i\},\\operatorname\{Pool\}\(V^\{\(l\)\}\)\_\{i\}\)
10:endfor

11:

Θ←Θ−η​∇Θ\(12​m​∑l=1m∑i=1nv\(κ1\(i,l\)\+κ2\(i,l\)\)\)\\Theta\\leftarrow\\Theta\-\\eta\\nabla\_\{\\Theta\}\\big\(\\frac\{1\}\{2m\}\\sum\_\{l=1\}^\{m\}\\sum\_\{i=1\}^\{n\_\{v\}\}\(\\kappa\_\{1\}^\{\(i,l\)\}\+\\kappa\_\{2\}^\{\(i,l\)\}\)\\big\)// Key\-value distillation

12:endfor

13:foreach

\(X,q\)\(X,q\)do

14:for

i∈\{1,…,g\}i\\in\\\{1,\\ldots,g\\\}do

15:

\{\(K~\(l\),V~\(l\)\)\}l∈Gi←fW\(i\)​\(hT\(i\)​\(X\)\)\\\{\(\\tilde\{K\}^\{\(l\)\},\\tilde\{V\}^\{\(l\)\}\)\\\}\_\{l\\in G\_\{i\}\}\\leftarrow f\_\{W\}^\{\(i\)\}\(h\_\{T\}^\{\(i\)\}\(X\)\)
16:endfor

17:

\(P1,\{ψ1\(l\)\}\)←f⁡\(\(K~,V~\),q\)\(P\_\{1\},\\\{\\psi\_\{1\}^\{\(l\)\}\\\}\)\\leftarrow f\(\(\\tilde\{K\},\\tilde\{V\}\),q\)
18:

\(P2,\{ψ2\(l\)\}\)←f⁡\(X,q\)\(P\_\{2\},\\\{\\psi\_\{2\}^\{\(l\)\}\\\}\)\\leftarrow f\(X,q\)
19:

L←λ1DKL\(P1∥P2\)\+λ2m∑l=1m\(1−cos\(ψ1\(l\),ψ2\(l\)\)\)L\\leftarrow\\lambda\_\{1\}D\_\{\\mathrm\{KL\}\}\(P\_\{1\}\\\|P\_\{2\}\)\+\\frac\{\\lambda\_\{2\}\}\{m\}\\sum\_\{l=1\}^\{m\}\(1\-\\cos\(\\psi\_\{1\}^\{\(l\)\},\\psi\_\{2\}^\{\(l\)\}\)\)
20:

Θ←Θ−η​∇ΘL\\Theta\\leftarrow\\Theta\-\\eta\\nabla\_\{\\Theta\}L// Output distillation

21:endfor

22:return

Θ^\\hat\{\\Theta\}

We empirically verify this on Qwen\-7B\-Instruct and Llama\-3\-8B\-Instruct models withG=4G\{=\}4layer groups andk=50k\{=\}50demonstrations\. The rapid coefficient decay in Figure[2](https://arxiv.org/html/2609.17888#S3.F2)indicates that a finite\-dimensional HiPPO state can retain most of the information in the KV cache\. Based on this observation, the SSMs process the entireTT\-token input using recurrent scans withO⁡\(T\)O\(T\)cost for a fixed number of groups\. The transformer then processes the resultingnvn\_\{v\}tokens with a cost ofO⁡\(nv2\)O\(n\_\{v\}^\{2\}\)\. Sincenvn\_\{v\}is fixed andnv≪Tn\_\{v\}\\ll T, the total prefix\-processing complexity isO⁡\(T\+nv2\)=O⁡\(T\)O\(T\+n\_\{v\}^\{2\}\)=O\(T\)\. In the next part, we introduce a two\-stage distillation method to train the SSMs to produce KV states that preserve the original transformer function output\.

Figure 2:Mean squared discrete Legendre coefficients of KV caches on Qwen\-7B\-Instruct and Llama\-3\-8B\-Instruct withG=4G\{=\}4layer groups andk=50k\{=\}50demonstrations\. The coefficients decay rapidly across groups, and the KV cache is dominated by low\-order HiPPO\-LegS components\.#### Key, value, and output distillation\.

To train the SSMs, we design a distillation procedure\. First, we align the generated tokens with a valid KV space\. Given themm\-layer modelfWf\_\{W\}, we extract its exact KV cache\{K\(l\),V\(l\)\}l=1m\\\{K^\{\(l\)\},V^\{\(l\)\}\\\}\_\{l=1\}^\{m\}\. To resolve the length mismatch, we apply pooling to theTTstates, producingnvn\_\{v\}pooled targetsPool⁡\(K\)\\operatorname\{Pool\}\(K\)andPool⁡\(V\)\\operatorname\{Pool\}\(V\)\. For theii\-th layer group, we keep the HiPPO\-LegS transition matrix fixed and train only the SSM input projectionB¯\(i\)\\bar\{B\}^\{\(i\)\}and the parameters of the KV projection MLPf\(i\)f^\{\(i\)\}\. We denote all trainable parameters byΘ=\{B¯\(i\),W\(i\)\}i=1g\\Theta=\\left\\\{\\bar\{B\}^\{\(i\)\},W^\{\(i\)\}\\right\\\}\_\{i=1\}^\{g\}\. The objective is:

12​m​∑l=1m∑i=1nv\(\(1−cos⁡\(K~i\(l\),Pool⁡\(K\(l\)\)i\)\)CLOSE\\displaystyle\\frac\{1\}\{2m\}\\sum\_\{l=1\}^\{m\}\\sum\_\{i=1\}^\{n\_\{v\}\}\\Big\(\\left\(1\-\\cos\\\!\\left\(\\tilde\{K\}^\{\(l\)\}\_\{i\},\\operatorname\{Pool\}\(K^\{\(l\)\}\)\_\{i\}\\right\)\\right\)OPEN\+\(1−cos⁡\(V~i\(l\),Pool⁡\(V\(l\)\)i\)\)\)\.\\displaystyle\\qquad\+\\left\(1\-\\cos\\\!\\left\(\\tilde\{V\}^\{\(l\)\}\_\{i\},\\operatorname\{Pool\}\(V^\{\(l\)\}\)\_\{i\}\\right\)\\right\)\\Big\)\.
Model\# DemonstrationsSST\-2Poem SentimentCoin FlipModular AdditionQwen\-1\.5Bk=50k=502\.5±0\.1×10−32\.5\_\{\\pm 0\.1\}\\times 10^\{\-3\}1\.3±0\.1×10−31\.3\_\{\\pm 0\.1\}\\times 10^\{\-3\}1\.9±0\.5×10−31\.9\_\{\\pm 0\.5\}\\times 10^\{\-3\}1\.3±0\.1×10−31\.3\_\{\\pm 0\.1\}\\times 10^\{\-3\}k=75k=752\.8±0\.2×10−32\.8\_\{\\pm 0\.2\}\\times 10^\{\-3\}1\.9±0\.2×10−31\.9\_\{\\pm 0\.2\}\\times 10^\{\-3\}2\.1±0\.2×10−32\.1\_\{\\pm 0\.2\}\\times 10^\{\-3\}1\.8±0\.4×10−31\.8\_\{\\pm 0\.4\}\\times 10^\{\-3\}Qwen\-3Bk=50k=506\.1±0\.2×10−36\.1\_\{\\pm 0\.2\}\\times 10^\{\-3\}2\.3±0\.3×10−32\.3\_\{\\pm 0\.3\}\\times 10^\{\-3\}2\.1±0\.8×10−32\.1\_\{\\pm 0\.8\}\\times 10^\{\-3\}2\.5±0\.3×10−32\.5\_\{\\pm 0\.3\}\\times 10^\{\-3\}k=75k=756\.3±0\.2×10−36\.3\_\{\\pm 0\.2\}\\times 10^\{\-3\}2\.9±0\.2×10−32\.9\_\{\\pm 0\.2\}\\times 10^\{\-3\}4\.1±1\.1×10−34\.1\_\{\\pm 1\.1\}\\times 10^\{\-3\}2\.1±0\.7×10−32\.1\_\{\\pm 0\.7\}\\times 10^\{\-3\}Llama\-8Bk=50k=502\.3±0\.1×10−32\.3\_\{\\pm 0\.1\}\\times 10^\{\-3\}2\.7±0\.1×10−32\.7\_\{\\pm 0\.1\}\\times 10^\{\-3\}1\.5±0\.5×10−31\.5\_\{\\pm 0\.5\}\\times 10^\{\-3\}3\.8±1\.0×10−33\.8\_\{\\pm 1\.0\}\\times 10^\{\-3\}k=75k=752\.5±0\.2×10−32\.5\_\{\\pm 0\.2\}\\times 10^\{\-3\}2\.9±0\.2×10−32\.9\_\{\\pm 0\.2\}\\times 10^\{\-3\}1\.9±0\.2×10−31\.9\_\{\\pm 0\.2\}\\times 10^\{\-3\}4\.8±0\.9×10−34\.8\_\{\\pm 0\.9\}\\times 10^\{\-3\}Table 1:Relative error of the predictive output between the full demonstration prefixppand the tokensp^\\hat\{p\}\. We vary the number of demonstrations on multiple tasks\. Lower values indicate higher fidelity to the original model output\. We run three times to compute the mean and standard deviations\.Next, we train the SSM using joint supervision of outputs and intermediate states to preserve predictive accuracy and positional structure\. For a queryqq, letPr\[y∣K~,V~,q\]\\Pr\[y\\mid\\tilde\{K\},\\tilde\{V\},q\]andPr\[y∣X,q\]\\Pr\[y\\mid X,q\]denote the output distributions of the LLM conditioned on the SSMs’ output and original demonstrations, respectively\. We denote these two distributions byP1P\_\{1\}andP2P\_\{2\}, respectively\. We additionally align the corresponding query hidden statesψ1\(l\)\\psi\_\{1\}^\{\(l\)\}andψ2\(l\)\\psi\_\{2\}^\{\(l\)\}across all layers\. The entire loss objective, denoted byL^​\(Θ\)\\hat\{L\}\(\\Theta\), is:

λ1DKL\(P1∥P2\)\+λ2m∑l=1m\(1−cos\(ψ1\(l\),ψ2\(l\)\)\)\.\\displaystyle\\lambda\_\{1\}D\_\{\\mathrm\{KL\}\}\\\!\\left\(P\_\{1\}\\\|P\_\{2\}\\right\)\+\\frac\{\\lambda\_\{2\}\}\{m\}\\sum\_\{l=1\}^\{m\}\\left\(1\-\\cos\\\!\\left\(\\psi\_\{1\}^\{\(l\)\},\\psi\_\{2\}^\{\(l\)\}\\right\)\\right\)\.whereλ1\\lambda\_\{1\}andλ2\\lambda\_\{2\}adjust the relative weight of the respective terms\. The entire procedure is summarized in Algorithm[1](https://arxiv.org/html/2609.17888#alg1)\.

To validate the effectiveness of the above distillation procedure, we evaluate Algorithm[1](https://arxiv.org/html/2609.17888#alg1)on multiple tasks usingkkdemonstrations from5050to7575\. As shown in Table[1](https://arxiv.org/html/2609.17888#S3.T1), the relative error between the logits from the full prefixppand tokensp^\\hat\{p\}remains strictly below0\.7%0\.7\\%across all settings\. This confirms that Algorithm[1](https://arxiv.org/html/2609.17888#alg1)robustly preserves the original predictive behavior\. We defer the omitted result wherek=25k=25to Table[7](https://arxiv.org/html/2609.17888#A2.T7)in Appendix[B\.3](https://arxiv.org/html/2609.17888#A2.SS3)\.

### 3\.2Demonstration Selection Using SSMs

We now utilize Algorithm[1](https://arxiv.org/html/2609.17888#alg1)for demonstration selection\. We first introduce a randomized estimator that quantifies the marginal utility of each candidate via subset sampling\. We define the affinity score,μ⁡\(z\)\\mu\(z\), of a demonstration examplez∈𝒟z\\in\\mathcal\{D\}as

−𝔼D′⊆D,z∈D′​\[𝔼\(q,y\)∼𝒫​\[ℓ⁡\(fW​\(D′,q\),y\)\]\]\.\-\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}D^\{\\prime\}\\subseteq\{D\},z\\in D^\{\\prime\}\\end\{subarray\}\}\[\\mathbb\{E\}\_\{\(q,y\)\\sim\\mathcal\{P\}\}\[\\ell\(f\_\{W\}\(D^\{\\prime\},q\),y\)\]\]\.Thus,μ⁡\(z\)\\mu\(z\)is the negative expected inference loss averaged over subsets of sizekkthat containzz\.

Algorithm 2Demonstration Selection Using SSMsInput: Candidate setDD, validation set𝒱\\mathcal\{V\}, subset sizekk, number of sampled subsetsMM, parameter initializationΘ\\Theta Require: Transformer modelff, groups\{Gi\}i=1g\\\{G\_\{i\}\\\}\_\{i=1\}^\{g\}, HiPPO\-LegS matrixA¯\\bar\{A\} Output:kkselected demonstrations

1:

Θ^←\\hat\{\\Theta\}\\leftarrowAlgorithm[1](https://arxiv.org/html/2609.17888#alg1)

2:for

j∈\{1,…,M\}j\\in\\\{1,\\ldots,M\\\}do

3:Sample

Dj′⊆DD\_\{j\}^\{\\prime\}\\subseteq Dof size

kk
4:

X←X\\leftarrowEmbeddings of

Dj′D\_\{j\}^\{\\prime\}
5:

T←T\\leftarrowLength of

XX
6:for

i∈\{1,…,g\}i\\in\\\{1,\\ldots,g\\\}do

7:

\{\(K~\(l\),V~\(l\)\)\}l∈Gi←fW\(i\)​\(hT​\(X\)\)\\\{\(\\tilde\{K\}^\{\(l\)\},\\tilde\{V\}^\{\(l\)\}\)\\\}\_\{l\\in G\_\{i\}\}\\leftarrow f\_\{W\}^\{\(i\)\}\(h\_\{T\}\(X\)\)
8:endfor

9:

C~←\{\(K~\(l\),V~\(l\)\)\}l=1m\\tilde\{C\}\\leftarrow\\\{\(\\tilde\{K\}^\{\(l\)\},\\tilde\{V\}^\{\(l\)\}\)\\\}\_\{l=1\}^\{m\}
10:

L^j←1\|𝒱\|​∑\(q,y\)∈𝒱ℓ⁡\(f⁡\(C~,q\),y\)\\displaystyle\\hat\{L\}\_\{j\}\\leftarrow\\frac\{1\}\{\|\\mathcal\{V\}\|\}\\sum\_\{\(q,y\)\\in\\mathcal\{V\}\}\\ell\\\!\\left\(f\(\\tilde\{C\},q\),y\\right\)
11:endfor

12:for

z∈Dz\\in Ddo

13:

μ^​\(z\)←−∑j=1M𝟏\{z∈D′j\}L^j∑j=1M𝟏\{z∈D′j\}\\displaystyle\\hat\{\\mu\}\(z\)\\leftarrow\-\\frac\{\\sum\_\{j=1\}^\{M\}\\mathbf\{1\}\_\{\\\{z\\in D^\{\\prime\}\_\{j\}\\\}\}\\hat\{L\}\_\{j\}\}\{\\sum\_\{j=1\}^\{M\}\\mathbf\{1\}\_\{\\\{z\\in D^\{\\prime\}\_\{j\}\\\}\}\}
14:endfor

15:returnTop\-

kkdemonstrations according to

μ^​\(z\)\\hat\{\\mu\}\(z\)for all

z∈Dz\\in D

To estimateμ⁡\(z\)\\mu\(z\), we uniformly sampleMMsubsets\{D1′,…,DM′\}\\\{D^\{\\prime\}\_\{1\},\\dots,D^\{\\prime\}\_\{M\}\\\}fromDDsizedkk\. Then, we estimate the affinity, denoted byμ^​\(z\)\\hat\{\\mu\}\(z\), as

−∑i=1M𝟏z∈Di′​∑\(q,y\)∈𝒱ℓ⁡\(fW​\(Di′,q\),y\)∑i=1M𝟏z∈Di′\.\-\\frac\{\\sum\_\{i=1\}^\{M\}\\mathbf\{1\}\_\{z\\in D^\{\\prime\}\_\{i\}\}\\sum\_\{\(q,y\)\\in\\mathcal\{V\}\}\\ell\(f\_\{W\}\(D^\{\\prime\}\_\{i\},q\),y\)\}\{\\sum\_\{i=1\}^\{M\}\\mathbf\{1\}\_\{z\\in D^\{\\prime\}\_\{i\}\}\}\.It is possible to show that withM=O⁡\(n/\(k​ϵ2\)\)M=O\(n/\(k\\epsilon^\{2\}\)\)sampled subsets, one can reduce the uniform estimation error down toϵ\\epsilon[Li et al\. \(2023b\)](https://arxiv.org/html/2609.17888#bib.bib21);[Li et al\. \(2023a\)](https://arxiv.org/html/2609.17888#bib.bib28)\. Afterwards, we rank the candidates byμ^​\(z\)\\hat\{\\mu\}\(z\)and select thekkdemonstrations with the highest scores\. Taken together, we summarize the entire procedure in Algorithm[2](https://arxiv.org/html/2609.17888#alg2)\.

#### A case study of linear functions\.

To evaluate our algorithm, we adopt the synthetic linear regression setting from[Garg et al\. \(2022\)](https://arxiv.org/html/2609.17888#bib.bib1)using a two\-layer transformer\. Both demonstrations and queries are generated viayi=⟨β\(c\),xi⟩\+ϵiy\_\{i\}=\\langle\\beta^\{\(c\)\},x\_\{i\}\\rangle\+\\epsilon\_\{i\}, where each inputxi∈ℝ20x\_\{i\}\\in\\mathbb\{R\}^\{20\}is independently sampled from the standard Gaussian distribution𝒩⁡\(0,Id20\)\\mathcal\{N\}\(0,\\id\_\{20\}\)\. To simulate multiple queries clustering around shared underlying tasks, each coefficient vectorβ\(c\)\\beta^\{\(c\)\}is drawn randomly from a set of mutually orthogonal anchors, ensuring strict task decoupling\. Finally, we inject independent Gaussian noiseϵi∼𝒩⁡\(0,σ2\)\\epsilon\_\{i\}\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}\)to prevent the model from trivially converging to an exact solution after observing2020examples\.

We evaluate the predictive performance of our proposed algorithm against two baseline strategies: random demonstration selection \(random\-kk\) and embedding similarity\-based selection \(top\-kk\)\. As shown in Figure[3\(a\)](https://arxiv.org/html/2609.17888#S3.F3.sf1), our method achieves the lowest estimation error in the linear setting\. This predictive advantage robustly extends to the non\-linear setting, which is a two\-layer ReLU network \(Figure[3\(b\)](https://arxiv.org/html/2609.17888#S3.F3.sf2)\)\. Note that the model is trained with4040data points as demonstrations, thus the error increases after the number of samples reaches4040\.

By leveraging the SSMs during subset evaluation, our approach reduces FLOP overhead by15\.7×15\.7\\timescompared to naive inference \(Figure[3](https://arxiv.org/html/2609.17888#S3.F3)\), ensuring high scalability for massive contexts\.

\(a\)Linear Functions\(b\)Nonlinear Functions\(c\)Linear Functions\(d\)Nonlinear Functions
Figure 3:Evaluation of our approach for linear and nonlinear functions\. The top panel provides a comparison with random\-kkand top\-kkselection methods\. The lower panel illustrates the computational costs relative to full inference \(denoted as not optimized\)\.MethodWindowMMLUModular AdditionGSM8KLengthErrorFLOPsMemoryErrorFLOPsMemoryErrorFLOPsMemoryDense\-\-6\.53​e156\.53e^\{15\}17\.3117\.31G\-1\.05​e151\.05e^\{15\}13\.6513\.65G\-2\.59​e152\.59e^\{15\}15\.7515\.75GStreamingLLM64641\.1%1\.1\\%7\.75​e137\.75e^\{13\}15\.9015\.90G9\.2%9\.2\\%1\.90​e131\.90e^\{13\}12\.9512\.95G5\.6%5\.6\\%1\.26​e141\.26e^\{14\}14\.0614\.06GLM\-Infinite64641\.2%1\.2\\%7\.75​e137\.75e^\{13\}15\.9115\.91G10\.9%10\.9\\%1\.91​e131\.91e^\{13\}12\.9712\.97G5\.1%5\.1\\%1\.27​e141\.27e^\{14\}14\.0614\.06GGist Tokens64647\.2%7\.2\\%1\.66​e151\.66e^\{15\}14\.4514\.45G2\.9%2\.9\\%7\.50​e147\.50e^\{14\}13\.7213\.72G40\.4%40\.4\\%2\.37​e142\.37e^\{14\}14\.4314\.43GICAE646416\.5%16\.5\\%3\.79​e143\.79e^\{14\}17\.0217\.02G1\.4%1\.4\\%1\.36​e141\.36e^\{14\}14\.2714\.27G0\.5%0\.5\\%3\.33​e143\.33e^\{14\}14\.8714\.87GDuoAttention64642\.7%2\.7\\%8\.03​e138\.03e^\{13\}16\.2416\.24G7\.7%7\.7\\%1\.93​e131\.93e^\{13\}13\.2513\.25G4\.9%4\.9\\%1\.31​e141\.31e^\{14\}14\.2514\.25GI2CL\-14\.3%14\.3\\%2\.62​e142\.62e^\{14\}15\.8915\.89G0\.7%0\.7\\%8\.29​e148\.29e^\{14\}13\.1413\.14G11\.7%11\.7\\%8\.56​e148\.56e^\{14\}13\.8613\.86GBSA∼400\\sim 4003\.1%3\.1\\%7\.86​e137\.86e^\{13\}16\.5316\.53G1\.8%1\.8\\%2\.31​e132\.31e^\{13\}13\.3613\.36G2\.1%2\.1\\%1\.51​e141\.51e^\{14\}14\.3114\.31GAlgorithm[1](https://arxiv.org/html/2609.17888#alg1)𝟏𝟔\\mathbf\{16\}0\.2%\\mathbf\{0\.2\\%\}3\.49​𝐞𝟏𝟑\\mathbf\{3\.49e^\{13\}\}16\.76\\mathbf\{16\.76\}G0\.2%\\mathbf\{0\.2\\%\}2\.12​𝐞𝟏𝟑\\mathbf\{2\.12e^\{13\}\}13\.27\\mathbf\{13\.27\}G0\.1%\\mathbf\{0\.1\\%\}8\.67​𝐞𝟏𝟑\\mathbf\{8\.67e^\{13\}\}14\.44\\mathbf\{14\.44\}G

Table 2:We report the relative error, computational cost \(FLOPs\), and memory usage on MMLU, Addition datasets, and GSM8K\. The relative error is computed against actual inference results\.CategorySentiment AnalysisMath ReasoningGraph ReasoningDatasetSST\-2Poem SentimentModular AdditionCoin FlipEdge ExistenceMetricAcc\.FLOPsAcc\.FLOPsAcc\.FLOPsAcc\.FLOPsAcc\.FLOPsRandom\-kk74\.7±4\.574\.7\_\{\\pm 4\.5\}1\.92​e141\.92e^\{14\}57\.3±4\.557\.3\_\{\\pm 4\.5\}4\.17​e144\.17e^\{14\}53\.2±1\.753\.2\_\{\\pm 1\.7\}1\.68​e141\.68e^\{14\}55\.3±7\.555\.3\_\{\\pm 7\.5\}2\.34​e142\.34e^\{14\}48\.9±3\.648\.9\_\{\\pm 3\.6\}6\.55​e146\.55e^\{14\}BM2577\.8±1\.977\.8\_\{\\pm 1\.9\}1\.92​e141\.92e^\{14\}55\.7±2\.855\.7\_\{\\pm 2\.8\}4\.17​e144\.17e^\{14\}58\.3±2\.558\.3\_\{\\pm 2\.5\}1\.68​e141\.68e^\{14\}42\.3±1\.142\.3\_\{\\pm 1\.1\}2\.34​e142\.34e^\{14\}51\.2±0\.651\.2\_\{\\pm 0\.6\}6\.55​e146\.55e^\{14\}Top\-kk88\.9±1\.288\.9\_\{\\pm 1\.2\}1\.92​e141\.92e^\{14\}63\.8±1\.563\.8\_\{\\pm 1\.5\}4\.17​e144\.17e^\{14\}60\.3±1\.160\.3\_\{\\pm 1\.1\}1\.68​e141\.68e^\{14\}45\.5±1\.445\.5\_\{\\pm 1\.4\}2\.34​e142\.34e^\{14\}58\.3±1\.158\.3\_\{\\pm 1\.1\}6\.55​e146\.55e^\{14\}Top\-kk\+ Alg\.[1](https://arxiv.org/html/2609.17888#alg1)88\.7±0\.888\.7\_\{\\pm 0\.8\}9\.14​e129\.14e^\{12\}63\.4±1\.063\.4\_\{\\pm 1\.0\}2\.17​e132\.17e^\{13\}60\.3±1\.160\.3\_\{\\pm 1\.1\}7\.51​e127\.51e^\{12\}45\.4±0\.845\.4\_\{\\pm 0\.8\}9\.56​e129\.56e^\{12\}58\.0±1\.358\.0\_\{\\pm 1\.3\}2\.72​e132\.72e^\{13\}GradCE83\.5±0\.783\.5\_\{\\pm 0\.7\}6\.71​e136\.71e^\{13\}70\.3±0\.970\.3\_\{\\pm 0\.9\}1\.47​e141\.47e^\{14\}64\.5±2\.164\.5\_\{\\pm 2\.1\}5\.80​e135\.80e^\{13\}68\.2±0\.868\.2\_\{\\pm 0\.8\}7\.80​e137\.80e^\{13\}65\.4±1\.365\.4\_\{\\pm 1\.3\}2\.61​e142\.61e^\{14\}BRIDGE89\.9±0\.389\.9\_\{\\pm 0\.3\}5\.73​e155\.73e^\{15\}73\.3±0\.873\.3\_\{\\pm 0\.8\}7\.12​e157\.12e^\{15\}70\.8±1\.370\.8\_\{\\pm 1\.3\}7\.01​e157\.01e^\{15\}82\.3±0\.682\.3\_\{\\pm 0\.6\}1\.00​e161\.00e^\{16\}76\.3±0\.476\.3\_\{\\pm 0\.4\}2\.84​e162\.84e^\{16\}BRIDGE \+ Alg\.[1](https://arxiv.org/html/2609.17888#alg1)89\.9±0\.589\.9\_\{\\pm 0\.5\}4\.93​e144\.93e^\{14\}73\.0±0\.973\.0\_\{\\pm 0\.9\}5\.70​e145\.70e^\{14\}70\.1±0\.570\.1\_\{\\pm 0\.5\}5\.49​e145\.49e^\{14\}82\.2±0\.882\.2\_\{\\pm 0\.8\}7\.22​e147\.22e^\{14\}75\.8±0\.875\.8\_\{\\pm 0\.8\}2\.08​e152\.08e^\{15\}Algorithm[2](https://arxiv.org/html/2609.17888#alg2)95\.9±0\.7\\mathbf\{95\.9\_\{\\pm 0\.7\}\}4\.52​𝐞𝟏𝟒\\mathbf\{4\.52e^\{14\}\}78\.6±1\.1\\mathbf\{78\.6\_\{\\pm 1\.1\}\}5\.23​𝐞𝟏𝟒\\mathbf\{5\.23e^\{14\}\}81\.8±1\.2\\mathbf\{81\.8\_\{\\pm 1\.2\}\}5\.08​𝐞𝟏𝟒\\mathbf\{5\.08e^\{14\}\}85\.3±1\.2\\mathbf\{85\.3\_\{\\pm 1\.2\}\}6\.61​𝐞𝟏𝟒\\mathbf\{6\.61e^\{14\}\}83\.6±0\.6\\mathbf\{83\.6\_\{\\pm 0\.6\}\}1\.81​𝐞𝟏𝟓\\mathbf\{1\.81e^\{15\}\}

Table 3:We report the test accuracy \(%\) and the computational cost \(FLOPs\) across five different kinds of datasets\. We compare our approach with several existing demonstration selection methods, using5050demonstrations\. We run each experiment with three random seeds to report the standard deviations\.

## 4Experiments

We evaluate our method by addressing the following three questions\. First, how accurately and efficiently does Algorithm[1](https://arxiv.org/html/2609.17888#alg1)approximate full\-context inference compared with existing methods? Second, how effective is Algorithm[2](https://arxiv.org/html/2609.17888#alg2)in terms of the accuracy and computational cost of downstream demonstration selection? Third, how well does the learned SSM scale to much longer contexts and transfer across different domains?

Through extensive experiments in both text and reasoning datasets, we provide positive answers to all three questions above\. Finally, we conduct detailed ablation studies to validate the design of our approach\.

### 4\.1Experiment Setup

Datasets\.We evaluate our approach across a diverse set of downstream tasks with different output formats\. SST\-2 and Poem Sentiment require categorical sentiment labels\. Coin Flip and Edge Existence require binary decisions\. Modular Addition requires numerical answers\.

We further evaluate the accuracy of Algorithm[1](https://arxiv.org/html/2609.17888#alg1)on MMLU and GSM8K\. MMLU covers knowledge\-intensive multiple\-choice questions across a broad range of academic and professional subjects, and GSM8K requires the model to generate multi\-step reasoning processes and final numerical answers\. Thus, our evaluation covers both short\-output classification tasks and longer\-form chain\-of\-thought reasoning tasks\.

#### Baselines\.

Regarding efficient inference methods, we consider full inference \(dense\), streaming LLM\([Xiao et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib6)\), LM\-infinite\([Han et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib11)\), Gist tokens\([Mu et al\., 2023](https://arxiv.org/html/2609.17888#bib.bib15)\), in\-context autoencoder \(ICAE\)\([Ge et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib14)\), block\-sparse attention \(BSA\)\([Xiao et al\., 2025a](https://arxiv.org/html/2609.17888#bib.bib12)\), and DuoAttention\([Xiao et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib13)\)as baselines\.

Then, we consider baselines that use different measures to rank the demonstrations\. These include random selection \(Random\-kk\), selection based on probabilistic relevance rankings \(BM25\), embedding similarities \(top\-kk\), Bayesian refinement, and iterative demonstration generation, for example \(BRIDGE\)\([Wan et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib10)\), cross\-entropy\-based selection with gradient estimation \(GradCE\)\([Zhang et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib2)\)\. We defer a detailed description of datasets, models, and baselines to Appendix[B\.1](https://arxiv.org/html/2609.17888#A2.SS1)\. All other implementation details are deferred to Appendix[B](https://arxiv.org/html/2609.17888#A2)\.

### 4\.2Experiment Results

#### Accuracy and efficiency of Algorithm[1](https://arxiv.org/html/2609.17888#alg1)\.

We evaluate our approach using the Qwen\-3B model with5050demonstrations for MMLU and Modular Addition and2020for GSM8K\. We compute the relative error of the output logits, the computational cost \(measured in the number of FLOPs\), and the peak memory usage for each baseline\.

As shown in Table[2](https://arxiv.org/html/2609.17888#S3.T2), Algorithm[1](https://arxiv.org/html/2609.17888#alg1)reduces the inference cost by approximately two orders of magnitude compared to Dense inference\. By compressing the global prefix into1616tokens, our method maintains a relative output error below0\.2%0\.2\\%across all tasks\. Compared to the baseline BSA, our method yields an average error reduction of2\.17%2\.17\\%and a35\.5%35\.5\\%decrease in FLOPs, with minimal memory overhead\. We defer the results on the remaining datasets and those using Qwen\-1\.5B to Table[10](https://arxiv.org/html/2609.17888#A2.T10)and Table[11](https://arxiv.org/html/2609.17888#A2.T11)\.

#### Accuracy and efficiency of Algorithm[2](https://arxiv.org/html/2609.17888#alg2)\.

We integrate our approach into existing selection frameworks to assess its compatibility and evaluate it across five datasets\. We evaluate the accuracy and FLOPs of the whole pipeline in Algorithm[2](https://arxiv.org/html/2609.17888#alg2)\.

Using the Qwen\-3B model with5050demonstrations, Table[3](https://arxiv.org/html/2609.17888#S3.T3)shows that our approach accelerates existing pipelines, with an average accuracy drop of0\.22%0\.22\\%across other baselines\. Compared to the baseline, BRIDGE, our approach yields an average accuracy improvement of6\.48%6\.48\\%while requiring much fewer FLOPs\. We analyze the efficiency of Algorithm[2](https://arxiv.org/html/2609.17888#alg2)in Appendix[B\.3](https://arxiv.org/html/2609.17888#A2.SS3)\.

#### Length and task generalization\.

To evaluate length generalization, we test the module, trained solely on5050demonstrations, on extended contexts containing up to10001000demonstrations\. As shown in Figure[4\(a\)](https://arxiv.org/html/2609.17888#S4.F4.sf1), our method maintains stable performance without additional fine\-tuning\. When increasing the number of demonstrations to10001000, the relative output error remains below1\.8%1\.8\\%across all datasets\.

\(a\)Length error\(b\)Runtime scaling
Figure 4:Figure[4\(a\)](https://arxiv.org/html/2609.17888#S4.F4.sf1)reports the relative output error as the number of demonstrationskkincreases\. Figure[4\(b\)](https://arxiv.org/html/2609.17888#S4.F4.sf2)reports the wall\-clock inference time askkincreases\.We further compare the wall\-clock time for processing a single query with full inference and our approach\. As shown in Table[4](https://arxiv.org/html/2609.17888#S4.T4)on the left, full inference retains the quadratic growth as the context becomes longer, while our approach scales approximately linearly\.

Finally, we evaluate cross\-task transfer across SST\-2, MMLU, Modular Addition, and GSM8K\. When adapting to a new task, we initialize it using the SSM trained on the source task and apply only the second\-stage fine\-tuning objective to the target task\.

Table[4](https://arxiv.org/html/2609.17888#S4.T4)on the right reports the relative error for each source\-target pair, where rows indicate the source training tasks and columns indicate the target evaluation tasks\. Adaptation takes65%65\\%fewer GPU hours than full training, and the relative error is lower than1\.0%1\.0\\%for all task pairs\.

### 4\.3Ablation Studies

We conduct ablation studies to examine the main components of our framework\.

#### Number of groups\.

First, we study the effects of the HiPPO state dimension and the number of SSM layer groups on SST\-2\. As shown in Table[4](https://arxiv.org/html/2609.17888#S4.T4), increasing either value reduces the test loss, but the improvement gradually becomes smaller\. In particular, the effect of increasing the HiPPO state dimension largely saturates at512512\.

#### Number of tokens\.

Then, we study the effect of the number of tokens,nvn\_\{v\}, using Qwen\-3B with5050demonstrations\. We varynvn\_\{v\}from22to3232and report the output mean squared error \(MSE\) and the accuracy drop relative to full\-context inference\. As shown in Table[5](https://arxiv.org/html/2609.17888#S4.T5), increasingnvn\_\{v\}generally reduces both metrics\. The results become stable atnv=16n\_\{v\}=16\. Thus, we usenv=16n\_\{v\}=16in the main experiments\.

Table 4:Left: cross\-domain transfer error between training and test tasks\. Right: test loss under different HiPPO state dimensions and numbers of layer groups\.
#### Empirical sample size scaling\.

Finally, we study the empirical sample complexity of Algorithm[2](https://arxiv.org/html/2609.17888#alg2)\. First, we study the empirical sample complexity of the randomized ensemble method using linear functions\([Garg et al\., 2022](https://arxiv.org/html/2609.17888#bib.bib1)\)\. Following the setting of[Garg et al\. \(2022\)](https://arxiv.org/html/2609.17888#bib.bib1), we construct a mixed candidate pool containing demonstrations from two distributions,𝒟1\\mathcal\{D\}\_\{1\}and𝒟2\\mathcal\{D\}\_\{2\}, while sampling all validation queries from𝒟1\\mathcal\{D\}\_\{1\}\. We vary both the number of demonstrations from𝒟1\\mathcal\{D\}\_\{1\}in each subset and the number of subsets sampled by the ensemble\. The results show that each subset requires more than2020demonstrations from𝒟1\\mathcal\{D\}\_\{1\}to achieve near\-zero validation loss\. Moreover, sampling more than5050subsets produces a clear separation between the average affinity scores assigned to demonstrations from𝒟1\\mathcal\{D\}\_\{1\}and𝒟2\\mathcal\{D\}\_\{2\}\. Table[6](https://arxiv.org/html/2609.17888#S4.T6)in Appendix[B\.3](https://arxiv.org/html/2609.17888#A2.SS3)reports the full results\.

DatasetMetricnv=2n\_\{v\}=2nv=4n\_\{v\}=4nv=16n\_\{v\}=16nv=32n\_\{v\}=32SST\-2MSE0\.0270\.0270\.0150\.0150\.0060\.0060\.0030\.003Accuracy drop0\.070\.070\.050\.050\.010\.010\.010\.01Poem SentimentMSE0\.0100\.0100\.0060\.0060\.0020\.0020\.0020\.002Accuracy drop0\.060\.060\.050\.050\.040\.040\.040\.04GSM8KMSE0\.0020\.0020\.0010\.0010\.0020\.0020\.0010\.001Accuracy drop0\.020\.020\.020\.020\.010\.010\.010\.01Table 5:Varying number of tokens using Qwen\-3B with5050demonstrations\. We report the output MSE and the accuracy drop relative to full\-context inference on SST\-2, Poem Sentiment, and GSM8K\.\# Subsets10105050100100200200Samples from𝒟1\\mathcal\{D\}\_\{1\}0\.22±0\.040\.22\_\{\\pm 0\.04\}0\.55±0\.030\.55\_\{\\pm 0\.03\}0\.57±0\.040\.57\_\{\\pm 0\.04\}0\.56±0\.030\.56\_\{\\pm 0\.03\}Samples from𝒟2\\mathcal\{D\}\_\{2\}0\.21±0\.010\.21\_\{\\pm 0\.01\}0\.52±0\.040\.52\_\{\\pm 0\.04\}0\.51±0\.060\.51\_\{\\pm 0\.06\}0\.48±0\.040\.48\_\{\\pm 0\.04\}Table 6:To illustrate the scoring mechanism in our random ensemble method, we report the average score of demonstration samples from𝒟1\\mathcal\{D\}\_\{1\}and𝒟2\\mathcal\{D\}\_\{2\}, given the query set generated from𝒟1\\mathcal\{D\}\_\{1\}\. We vary the number of subsets and find that the in\-distribution samples always achieve a higher score than out\-of\-distribution samples\.

## 5Related Work

Long\-context model inference\.Work on reducing transformer cost can be grouped into four lines\. First, architectural methods use SSMs\([Gu et al\., 2020](https://arxiv.org/html/2609.17888#bib.bib5);[Gu et al\., 2022](https://arxiv.org/html/2609.17888#bib.bib4)\)and hybrid mamba\-transformer designs\([Gu and Dao, 2024](https://arxiv.org/html/2609.17888#bib.bib9);[Ren et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib8);[Oncescu et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib7)\)to achieve linear\-time long\-sequence modeling\. For long\-convolution sequence models, the fast Fourier transform\([Oncescu et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib7)\)can be utilized to perform exact autoregressive inference efficiently\. Note that these approaches target alternative architectures and cannot be directly applied to transformer\-based generative models\. The second line of work\([Xiao et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib6);[Han et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib11)\)limits each query to a sliding window of recent tokens and a few fixed initial tokens that act as attention sinks at inference time\. Third, memory\-token methods map long contexts into learned slots\([Mu et al\., 2023](https://arxiv.org/html/2609.17888#bib.bib15);[Chevalier et al\., 2023](https://arxiv.org/html/2609.17888#bib.bib16);[Ge et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib14)\), but they still need to run a transformer over the original input\. Finally, previous works[Yang et al\. \(2024\)](https://arxiv.org/html/2609.17888#bib.bib29);[Yang et al\. \(2026\)](https://arxiv.org/html/2609.17888#bib.bib30);[Dao et al\. \(2022\)](https://arxiv.org/html/2609.17888#bib.bib31)use blockwise computation to improve hardware efficiency for long sequences\. Given their insights, we apply the same general strategy to the fixed HiPPO recurrence by combining the state updates within each chunk into a single batched operation\.

#### Demonstration selection and prompt optimization\.

In\-context learning relies heavily on the quality and structure of the provided demonstrations\([Min et al\., 2022](https://arxiv.org/html/2609.17888#bib.bib3)\)\. Traditional selection relies on static embedding retrieval, which ignores the model’s actual output\. To capture these combinatorial effects, recent methods incorporate model inference feedback\. For instance,[Wan et al\. \(2025\)](https://arxiv.org/html/2609.17888#bib.bib10)alternate between selecting influential demonstrations and using them to regenerate an improved many\-shot prompt\. Since running the full model on every candidate subset is expensive,[Zhang et al\. \(2025b\)](https://arxiv.org/html/2609.17888#bib.bib2)estimate subset\-level outputs from input gradients and reduce the cost of evaluating many subsets\.

Another related line of work involves surrogate modeling[Li et al\. \(2023b\)](https://arxiv.org/html/2609.17888#bib.bib21), task attribution[Zhang et al\. \(2025a\)](https://arxiv.org/html/2609.17888#bib.bib27);[Zhang et al\. \(2026\)](https://arxiv.org/html/2609.17888#bib.bib25), and influence estimation\. A theoretical analysis\([Yang et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib32)\)shows that the benefit of combining datasets depends jointly on task mismatch and source sample size\. Adding more source data can even turn positive transfer into negative transfer\.[Li et al\. \(2023b\)](https://arxiv.org/html/2609.17888#bib.bib21);[Li et al\. \(2024a\)](https://arxiv.org/html/2609.17888#bib.bib18)use surrogate modeling to efficiently identify negative transfer among data, and[Zhang et al\. \(2026\)](https://arxiv.org/html/2609.17888#bib.bib25)use kernel surrogate modeling to estimate affinity\. The same goal of estimating interactions among datasets has also been studied for large language models\. Gradient\-based methods predict the fine\-tuning performance of different dataset combinations without training a separate model for each combination\([Li et al\., 2024b](https://arxiv.org/html/2609.17888#bib.bib19)\), and these estimates can then guide dataset grouping\([Li et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib20)\)\. Broadly, estimated task affinity can guide not only dataset grouping but also how model components are shared across tasks\([Li et al\., 2025a](https://arxiv.org/html/2609.17888#bib.bib26)\)\.

Beyond gradient\-based estimation and heuristic scoring, selecting efficient candidate data points connects to the broader literature on experimental design and automated trajectory optimization from historical data\([Xiong et al\., 2024a](https://arxiv.org/html/2609.17888#bib.bib24);[Xiong et al\., 2024b](https://arxiv.org/html/2609.17888#bib.bib23);[Xiong, 2025](https://arxiv.org/html/2609.17888#bib.bib22);[Zhang et al\., 2025a](https://arxiv.org/html/2609.17888#bib.bib27)\)\. It would be interesting to connect ideas from this paper to automated trajectory optimization from historical data\.

## 6Conclusion

We address the computational bottleneck of demonstration selection with SSMs, where many candidate subsets require repeated transformer inference\. We train the SSMs to estimate the transformer output for different demonstration subsets\. It reduce the inference complexity to linear\. Then, we introduce a random selection method that compares many subsets and identifies useful demonstrations based on the SSMs inference output\.

## Acknowledgment

Thanks to the anonymous reviewers and the area chair for their constructive feedback\. The work of Ziniu Zhang and Zhenshuo Zhang is partially funded by NSF award IIS\-2412008, a startup fund from Northeastern University, and Khoury PhD fellowships\.

## References

- Chevalieret al\.\(2023\)A\. Chevalier, A\. Wettig, A\. Ajith, and D\. ChenAdapting language models to compress contexts\.InConference on Empirical Methods in Natural Language Processing \(EMNLP\),pp\. 3829–3846\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Daoet al\.\(2022\)T\. Dao, D\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashattention: fast and memory\-efficient exact attention with io\-awareness\.Advances in Neural Information Processing Systems \(NeurIPS\)35,pp\. 16344–16359\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Garget al\.\(2022\)S\. Garg, D\. Tsipras, P\. S\. Liang, and G\. ValiantWhat can transformers learn in\-context? a case study of simple function classes\.Advances in Neural Information Processing Systems \(NeurIPS\)35,pp\. 30583–30598\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p1.1),[§3\.2](https://arxiv.org/html/2609.17888#S3.SS2.SSS0.Px1.p1.1),[§4\.3](https://arxiv.org/html/2609.17888#S4.SS3.SSS0.Px3.p1.1)\.
- Geet al\.\(2024\)T\. Ge, H\. Jing, L\. Wang, X\. Wang, S\. Chen, and F\. WeiIn\-context autoencoder for context compression in a large language model\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p16.1),[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Guet al\.\(2020\)A\. Gu, T\. Dao, S\. Ermon, A\. Rudra, and C\. ReHiPPO: recurrent memory with optimal polynomial projections\.Advances in Neural Information Processing Systems \(NeurIPS\)33\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§2](https://arxiv.org/html/2609.17888#S2.p4.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Gu and Dao \(2024\)A\. Gu and T\. DaoMamba: linear\-time sequence modeling with selective state spaces\.InFirst conference on language modeling \(COLM\),Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§2](https://arxiv.org/html/2609.17888#S2.p4.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Guet al\.\(2022\)A\. Gu, K\. Goel, and C\. ReEfficiently modeling long sequences with structured state spaces\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§2](https://arxiv.org/html/2609.17888#S2.p3.2),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Hanet al\.\(2024\)C\. Han, Q\. Wang, H\. Peng, W\. Xiong, Y\. Chen, H\. Ji, and S\. WangLm\-infinite: zero\-shot extreme length generalization for large language models\.InConference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(NAACL\-HLT\),pp\. 3991–4008\.Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p14.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Liet al\.\(2023a\)D\. Li, H\. Ju, A\. Sharma, and H\. R\. ZhangBoosting multitask learning on graphs through higher\-order task affinities\.InProceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining,pp\. 1213–1222\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p6.1),[§3\.2](https://arxiv.org/html/2609.17888#S3.SS2.p2.2)\.
- Liet al\.\(2023b\)D\. Li, H\. L\. Nguyen, and H\. R\. ZhangIdentification of negative transfers in multitask learning using surrogate models\.Transactions on Machine Learning Research \(TMLR\)\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p6.1),[§3\.2](https://arxiv.org/html/2609.17888#S3.SS2.p2.2),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Liet al\.\(2024a\)D\. Li, A\. Sharma, and H\. R\. ZhangScalable multitask learning using gradient\-based estimation of task affinity\.InACM SIGKDD Conference on Knowledge Discovery and Data Mining \(KDD\),pp\. 1542–1553\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p6.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Liet al\.\(2025a\)D\. Li, Z\. Zhang, M\. Duan, E\. Dobriban, and H\. R\. ZhangEfficiently learning branching networks for multitask algorithmic reasoning\.arXiv preprint arXiv:2512\.01113\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Liet al\.\(2024b\)D\. Li, Z\. Zhang, L\. Wang, and H\. R\. ZhangScalable fine\-tuning from multiple data sources: a first\-order approximation approach\.InFindings of the Association for Computational Linguistics: EMNLP 2024,pp\. 5608–5623\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p6.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Liet al\.\(2025b\)D\. Li, Z\. Zhang, L\. Wang, and H\. R\. ZhangEfficient ensemble for fine\-tuning language models on multiple datasets\.InAnnual Meeting of the Association for Computational Linguistics \(ACL\),pp\. 25347–25364\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Liet al\.\(2025c\)Z\. Li, Z\. Xu, L\. Han, Y\. Gao, S\. Wen, D\. Liu, H\. Wang, and D\. MetaxasImplicit in\-context learning\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p17.1)\.
- Minet al\.\(2022\)S\. Min, X\. Lyu, A\. Holtzman, M\. Artetxe, M\. Lewis, H\. Hajishirzi, and L\. ZettlemoyerRethinking the role of demonstrations: what makes in\-context learning work?\.InEmpirical Methods in Natural Language Processing \(EMNLP\),pp\. 11048–11064\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p1.1)\.
- Muet al\.\(2023\)J\. Mu, X\. Li, and N\. GoodmanLearning to compress prompts with gist tokens\.Advances in Neural Information Processing Systems \(NeurIPS\)36,pp\. 19327–19352\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Oncescuet al\.\(2025\)C\. Oncescu, S\. Purandare, S\. Idreos, and S\. M\. KakadeFlash inference: near linear time inference for long convolution sequence models and beyond\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p1.1),[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Renet al\.\(2025\)L\. Ren, Y\. Liu, Y\. Lu, Y\. Shen, C\. Liang, and W\. ChenSamba: simple hybrid state space models for efficient unlimited context language modeling\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Wanet al\.\(2025\)X\. Wan, H\. Zhou, R\. Sun, and S\. O\. ArikFrom few to many: self\-improving many\-shot reasoners through iterative optimization and generation\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p18.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p1.1)\.
- Xiaoet al\.\(2025a\)E\. Xiao, C\. Li, Y\. Zhang, G\. Neubig, and A\. BertschEfficient many\-shot in\-context learning with dynamic block\-sparse attention\.InAnnual Meeting of the Association for Computational Linguistics \(ACL\),pp\. 31946–31958\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1)\.
- Xiaoet al\.\(2025b\)G\. Xiao, J\. Tang, J\. Zuo, J\. Guo, S\. Yang, H\. Tang, Y\. Fu, and S\. HanDuoAttention: efficient long\-context LLM inference with retrieval and streaming heads\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p15.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1)\.
- Xiaoet al\.\(2024\)G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. LewisEfficient streaming language models with attention sinks\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p13.1),[§B\.2](https://arxiv.org/html/2609.17888#A2.SS2.p1.1),[§1](https://arxiv.org/html/2609.17888#S1.p3.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Xionget al\.\(2024a\)R\. Xiong, S\. Athey, M\. Bayati, and G\. ImbensOptimal experimental design for staggered rollouts\.Management Science70\(8\),pp\. 5317–5336\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p3.1)\.
- Xionget al\.\(2024b\)R\. Xiong, A\. Chin, and S\. J\. TaylorData\-driven switchback experiments: theoretical tradeoffs and empirical bayes designs\.arXiv preprint arXiv:2406\.06768\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p3.1)\.
- Xiong \(2025\)R\. XiongAutomated experimental design with optimization from historical data simulations\.Available at SSRN 5126080\.Cited by:[§1](https://arxiv.org/html/2609.17888#S1.p1.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p3.1)\.
- Yanget al\.\(2025\)F\. Yang, H\. R\. Zhang, S\. Wu, C\. Ré, and W\. J\. SuPrecise high\-dimensional asymptotics for quantifying heterogeneous transfers\.Journal of Machine Learning Research \(JMLR\)26\(113\),pp\. 1–88\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Yanget al\.\(2026\)S\. Yang, Y\. Shen, K\. Wen, S\. Tan, M\. Mishra, L\. Ren, R\. Panda, and Y\. KimPath attention: position encoding via accumulating householder transformations\.Advances in Neural Information Processing Systems \(NeurIPS\)38,pp\. 62220–62247\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Yanget al\.\(2024\)S\. Yang, B\. Wang, Y\. Zhang, Y\. Shen, and Y\. KimParallelizing linear transformers with the delta rule over sequence length\.Advances in Neural Information Processing Systems \(NeurIPS\)37,pp\. 115491–115522\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.p1.1)\.
- Zhanget al\.\(2025a\)Z\. Zhang, M\. Duan, Y\. Ye, and H\. R\. ZhangScalable multi\-objective and meta reinforcement learning via gradient estimation\.arXiv preprint arXiv:2511\.12779\.Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p3.1)\.
- Zhanget al\.\(2026\)Z\. Zhang, M\. Duan, and H\. R\. ZhangEfficient estimation of kernel surrogate models for task attribution\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p2.1)\.
- Zhanget al\.\(2025b\)Z\. Zhang, Z\. Zhang, D\. Li, L\. Wang, J\. Dy, and H\. R\. ZhangLinear\-time demonstration selection for in\-context learning via gradient estimation\.InEmpirical Methods in Natural Language Processing \(EMNLP\),Cited by:[§B\.1](https://arxiv.org/html/2609.17888#A2.SS1.p18.1),[§1](https://arxiv.org/html/2609.17888#S1.p1.1),[§1](https://arxiv.org/html/2609.17888#S1.p2.1),[§1](https://arxiv.org/html/2609.17888#S1.p6.1),[§4\.1](https://arxiv.org/html/2609.17888#S4.SS1.SSS0.Px1.p2.1),[§5](https://arxiv.org/html/2609.17888#S5.SS0.SSS0.Px1.p1.1)\.

## Appendix AProof of Proposition[3\.1](https://arxiv.org/html/2609.17888#S3.Thmtheorem1)

###### Proof\.

For completeness, the HiPPO\-LegS transition matrix used by each layer group is defined as

A¯p,q=\{−\(2​p\+1\)​\(2​q\+1\),if​p\>q,−\(p\+1\),if​p=q,0,if​p<q,\\bar\{A\}\_\{p,q\}=\\begin\{cases\}\-\\sqrt\{\(2p\+1\)\(2q\+1\)\},&\\text\{if \}p\>q,\\\\ \-\(p\+1\),&\\text\{if \}p=q,\\\\ 0,&\\text\{if \}p<q,\\end\{cases\}\(2\)
Let\{ϕr\}r=1T⊂ℝT\\\{\\phi\_\{r\}\\\}\_\{r=1\}^\{T\}\\subset\\mathbb\{R\}^\{T\}denote an orthonormal discrete Legendre basis for the token dimension, ordered by increasing polynomial degree, so

ϕr⊤ϕr′=𝟙\{r=r′\}\.\\phi\_\{r\}^\{\\top\}\\phi\_\{r^\{\\prime\}\}=\\mathbbm\{1\}\\\{r=r^\{\\prime\}\\\}\.For each layer groupGiG\_\{i\}, writeC\(i\)∈ℝT×DiC^\{\(i\)\}\\in\\mathbb\{R\}^\{T\\times D\_\{i\}\}, where the rows correspond to token positions and the columns contain the key and value states from all layers and key\-value heads in the group\. The vector of therr\-th HiPPO\-LegS coefficients is

cr\(i\):=\(C\(i\)\)⊤​ϕr∈ℝDi\.c\_\{r\}^\{\(i\)\}:=\\bigl\(C^\{\(i\)\}\\bigr\)^\{\\top\}\\phi\_\{r\}\\in\\mathbb\{R\}^\{D\_\{i\}\}\.Therefore,

C\(i\)=∑r=1Tϕr​\(cr\(i\)\)⊤\.C^\{\(i\)\}=\\sum\_\{r=1\}^\{T\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}\.SinceΠN\\Pi\_\{N\}projects along the token dimension onto the firstNNbasis functions,ΠN​C\(i\)=∑r=1Nϕr​\(cr\(i\)\)⊤\\Pi\_\{N\}C^\{\(i\)\}=\\sum\_\{r=1\}^\{N\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}\. The omitted part of the cache is thus

C\(i\)−ΠN​C\(i\)=∑r=N\+1Tϕr​\(cr\(i\)\)⊤\.C^\{\(i\)\}\-\\Pi\_\{N\}C^\{\(i\)\}=\\sum\_\{r=N\+1\}^\{T\}\\phi\_\{r\}\\bigl\(c\_\{r\}^\{\(i\)\}\\bigr\)^\{\\top\}\.By orthonormality of the basis,

‖C\(i\)−ΠN​C\(i\)‖F2=∑r=N\+1T‖cr\(i\)‖22\.\\left\\\|C^\{\(i\)\}\-\\Pi\_\{N\}C^\{\(i\)\}\\right\\\|\_\{F\}^\{2\}=\\sum\_\{r=N\+1\}^\{T\}\\left\\\|c\_\{r\}^\{\(i\)\}\\right\\\|\_\{2\}^\{2\}\.\(3\)Applying this inequality to each term in \([3](https://arxiv.org/html/2609.17888#A1.E3)\) gives

‖C\(i\)−ΠN​C\(i\)‖F2\\displaystyle\\left\\\|C^\{\(i\)\}\-\\Pi\_\{N\}C^\{\(i\)\}\\right\\\|\_\{F\}^\{2\}≤\\displaystyle\\leq1\(N\+1\)2​s​∑r=N\+1Tr2​s​‖cr\(i\)‖22,\\displaystyle\\frac\{1\}\{\(N\+1\)^\{2s\}\}\\sum\_\{r=N\+1\}^\{T\}r^\{2s\}\\left\\\|c\_\{r\}^\{\(i\)\}\\right\\\|\_\{2\}^\{2\},sincer2​s\(N\+1\)2​s≥1\\frac\{r^\{2s\}\}\{\(N\+1\)^\{2s\}\}\\geq 1forr≥N\+1r\\geq N\+1ands\>0s\>0\. ∎

## Appendix BOmitted Experiments

### B\.1Implementation Setup

Datasets\.We evaluate our method on datasets covering several task types, including language understanding, algorithmic reasoning, arithmetic reasoning, and graph reasoning\. The number of demonstrations within each task is15001500\.

The[SST\-2](https://huggingface.co/datasets/nyu-mll/glue/viewer/sst2)dataset is a binary sentiment classification benchmark from[GLUE](https://huggingface.co/datasets/nyu-mll/glue), consisting of movie reviews labeled as positive or negative\.

The[Poem Sentiment](https://huggingface.co/datasets/google-research-datasets/poem_sentiment)dataset contains sentiment labels for lines of poetry, with three classes: positive, neutral, and negative\. This dataset is used to evaluate sentiment understanding in literary text\.

Modular Addition is an arithmetic reasoning task in which the model is given two integersa,b∈\{0,…,P−1\}a,b\\in\\\{0,\\ldots,P\-1\\\}for a primePPand is asked to predict their sum moduloPP, that is, to predictccsuch thata\+b≡c\(modP\)a\+b\\equiv c\\pmod\{P\}\.

The[Coin Flip](https://huggingface.co/datasets/skrishna/coin_flip)dataset is a reasoning task in which the model is given a natural language description of a sequence of fair coin flips and must predict the final outcome, either heads or tails\.

The Edge Existence dataset from[GraphQA](https://github.com/google-research/google-research/tree/master/graphqa)is a graph reasoning task in which the model is given an undirected graph and must determine whether a specified edge is present in the graph\.

The[MMLU](https://huggingface.co/datasets/cais/mmlu)dataset is a multitask benchmark covering a wide range of subjects, including humanities, social sciences, and STEM\. It evaluates a model’s ability to perform multiple\-choice question answering across diverse domains\. We use the college\-level tasks in the evaluation\.

The[GSM8K](https://huggingface.co/datasets/openai/gsm8k)dataset is a grade\-school math reasoning benchmark consisting of natural language word problems\. Each example requires multi\-step arithmetic reasoning, and the model is asked to produce the final numerical answer\. We use this dataset to evaluate whether Algorithm[1](https://arxiv.org/html/2609.17888#alg1)preserves the reasoning information needed for long\-context mathematical inference\.

Models\.Our method is a model\-agnostic method\. In our experiments, we evaluate it on a diverse set of open\-source language models that vary in scale and architecture, including[Qwen2\.5\-1\.5B\-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct),[Qwen2\.5\-3B\-Instruct](https://huggingface.co/Qwen/Qwen2.5-3B-Instruct),[Qwen2\.5\-7B\-Instruct](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct), and[Llama\-3\.1\-8B\-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)\. This setup allows us to assess the effectiveness and scalability of our method across models with different capacities, training paradigms, and usage settings\.

Training budget\.We train our method on a single NVIDIA RTX A6000 GPU\. We set the number of virtual tokens to1616, the SSM state dimension to512512, and the number of layer groups to44\. Training for11epoch in phase one and22epochs in phase two requires1\.151\.15GPU hours\.

Baselines\.We compare Algorithm[1](https://arxiv.org/html/2609.17888#alg1)against full inference, StreamingLLM, LM\-Infinite, Block Sparse Attention baselines, and DuoAttention\. Then, we compare Algorithm[2](https://arxiv.org/html/2609.17888#alg2)against baselines including BM25, top\-kk, BRIDGE,GradCE\.

Full Attention serves as our baseline\. Every token attends to all preceding tokens, retaining the complete Key\-Value cache across all layers and heads\. While preserving maximum model capability, its linear memory scaling and quadratic pre\-filling cost become prohibitively expensive for large demonstration pools\.

StreamingLLM\([Xiao et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib6)\)exploits the phenomenon of attention sinks, where initial tokens receive disproportionately high attention scores regardless of semantic relevance\. It maintains a fixed\-size KV cache comprising only these initial sink tokens and a sliding window of recent tokens, evicting all intermediate context\. This achieves constant memory usage but inherently degrades performance on tasks requiring fact retrieval from the middle of the input\.

LM\-Infinite\([Han et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib11)\)addresses length generalization beyond the model’s pre\-training window\. To prevent attention degradation from out\-of\-distribution position indices, it discards tokens outside a prescribed local window and applies aΛ\\Lambda\-shaped position remapping\. This aligns the remaining tokens’ indices with the training distribution, enabling infinite context processing without fine\-tuning\. However, similar to StreamingLLM, it suffers from information loss by permanently evicting middle\-context tokens\.

DuoAttention\([Xiao et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib13)\)categorizes attention heads into retrieval heads \(attending to arbitrary positions\) and streaming heads \(attending to sinks and recent tokens\)\. During a lightweight identification phase, a trainable gateαi,j\\alpha\_\{i,j\}is assigned to each KV head and optimized via synthetic passkey\-retrieval data\. This minimizes output deviation from Full Attention under anℓ1\\ell\_\{1\}sparsity penalty\. At deployment, binarized gates selectively apply the StreamingLLM eviction policy exclusively to streaming heads, while retrieval heads retain the complete KV cache\. This selectively reduces memory and latency while preserving long\-range retrieval capabilities\.

ICAE\([Ge et al\., 2024](https://arxiv.org/html/2609.17888#bib.bib14)\)compresses the input context into a small number of learned memory tokens through an autoencoding objective\. It first encodes the original context into compact latent representations, and then conditions the language model on these memory tokens to recover the information needed for downstream generation\.

Implicit In\-context Learning \(I2CL\)\([Li et al\., 2025c](https://arxiv.org/html/2609.17888#bib.bib17)\)moves demonstrations from the token space into the model’s activation space\. It extracts a vector from each demonstration, combines these vectors in a way that is independent of their order, and injects the resulting context vector into the model’s residual streams during inference\. I2CL provides few\-shot performance with inference cost close to that of zero\-shot prediction\. However, constructing its context vector still requires Transformer forward passes over the demonstrations\.

The top\-kkmethod selects thekkmost similar candidates based on feature similarity, given a test input\. We compute the cosine similarity in the top\-kkrange between the last\-layer hidden representations of queries and demonstration examples\. The BM25 method utilizes a term frequency\-based ranking function to retrieve the top\-kkcandidates whose input texts are most relevant to the query, with a focus on lexical overlap rather than embedding\-level similarity\.GradCE\([Zhang et al\., 2025b](https://arxiv.org/html/2609.17888#bib.bib2)\)first filters candidates by embedding similarity, then selects a random anchor and performs a single forward and backward pass on all validation samples to precompute the loss and gradient with respect to the input embedding at that anchor\. For each remaining candidate, it computes the embedding difference from the anchor and takes the inner product with the precomputed gradient to approximate the candidate’s validation loss via a first\-order Taylor expansion, selecting thekkcandidates with the lowest estimated loss\. BRIDGE\([Wan et al\., 2025](https://arxiv.org/html/2609.17888#bib.bib10)\)iteratively alternates between using Bayesian optimization to select a small high\-performing subset of ICL demonstrations from the candidate pool \("optimize"\) and using that subset as seed examples to re\-generate reasoning paths on the training set back to the many\-shot regime \("generate"\), repeating this process over multiple rounds to progressively improve demonstration quality\.

Model\# DemonstrationsSST\-2CRPoem SentimentModular AdditionQwen\-1\.5Bk=25k=252\.3±0\.2×10−32\.3\_\{\\pm 0\.2\}\\times 10^\{\-3\}1\.9±0\.3×10−31\.9\_\{\\pm 0\.3\}\\times 10^\{\-3\}8\.9±2\.5×10−48\.9\_\{\\pm 2\.5\}\\times 10^\{\-4\}1\.1±0\.1×10−31\.1\_\{\\pm 0\.1\}\\times 10^\{\-3\}Qwen\-3Bk=25k=255\.3±0\.2×10−35\.3\_\{\\pm 0\.2\}\\times 10^\{\-3\}3\.2±0\.3×10−33\.2\_\{\\pm 0\.3\}\\times 10^\{\-3\}2\.0±2\.5×10−32\.0\_\{\\pm 2\.5\}\\times 10^\{\-3\}3\.5±0\.1×10−33\.5\_\{\\pm 0\.1\}\\times 10^\{\-3\}Table 7:Relative error of the predictive output between the full demonstration prefixppand the compressed virtual tokensp^\\hat\{p\}\. Evaluations are reported across multiple downstream tasks for varying demonstration counts\. Lower values indicate higher fidelity to the original uncompressed model output\. We run three times to compute the mean and standard deviations\.
### B\.2Efficient Implementation

Previous work[Xiao et al\. \(2024\)](https://arxiv.org/html/2609.17888#bib.bib6)shows that standard LLM inherently relies on initial tokens as “attention sinks” to stabilize the generation process\. Thus, we explicitly retain a small number of exact initial KV states as sink tokens from the raw prefix\. These states are concatenated with our generated virtual tokens to form the compressed context\.

The standard SSM recurrence processes tokens sequentially:

st=A​st−1\+ut,ut=B​xt,s0=0,s\_\{t\}=As\_\{t\-1\}\+u\_\{t\},u\_\{t\}=Bx\_\{t\},s\_\{0\}=0,\(4\)wherext∈ℝHx\_\{t\}\\in\\mathbb\{R\}^\{H\}is the input embedding at steptt,B∈ℝD×HB\\in\\mathbb\{R\}^\{D\\times H\}is a learned input projection,A∈ℝD×DA\\in\\mathbb\{R\}^\{D\\times D\}is the frozen discretized HiPPO matrix, andst∈ℝDs\_\{t\}\\in\\mathbb\{R\}^\{D\}is the hidden state\. A naive implementation requiresTTsequential kernel launches, which may under utilize the GPU\.

We observe that unrolling the recurrence overCCconsecutive steps yields a closed\-form expression\. Starting from an incoming statesins\_\{\\mathrm\{in\}\}, the state after processing inputsu0,u1,…,uC−1u\_\{0\},u\_\{1\},\\ldots,u\_\{C\-1\}is:

sout=AC​sin\+∑t=0C−1AC−1−t​ut\.s\_\{\\mathrm\{out\}\}=A^\{C\}s\_\{\\mathrm\{in\}\}\+\\sum\_\{t=0\}^\{C\-1\}A^\{C\-1\-t\}u\_\{t\}\.\(5\)The summation can be rewritten as a single batched contraction\. Define the weight matricesWt=AC−1−tW\_\{t\}=A^\{C\-1\-t\}fort=0,…,C−1t=0,\\ldots,C\{\-\}1, so thatW0=AC−1W\_\{0\}=A^\{C\-1\}\(oldest input, most decay\) andWC−1=IW\_\{C\-1\}=I\(newest input, no decay\)\. Stacking these into a tensor𝐖∈ℝC×D×D\\mathbf\{W\}\\in\\mathbb\{R\}^\{C\\times D\\times D\}and the chunk inputs into𝐔∈ℝC×D\\mathbf\{U\}\\in\\mathbb\{R\}^\{C\\times D\}, the entire summation reduces to:

∑t=0C−1Wt​ut=∑t=0C−1𝐖⁡\[t\]​𝐔​\[t\],\\sum\_\{t=0\}^\{C\-1\}W\_\{t\}u\_\{t\}=\\sum\_\{t=0\}^\{C\-1\}\\mathbf\{W\}\[t\]\\mathbf\{U\}\[t\],\(6\)which is evaluated as a singleeinsumoperation on the GPU\.

We partition the full prefix of lengthTTinto⌈T/C⌉\\lceil T/C\\rceilcontiguous chunks\. Processing proceeds as follows: for each chunkk=0,1,…k=0,1,\\ldots, we first compute the powers\{AC−1,AC−2,…,I\}\\\{A^\{C\-1\},A^\{C\-2\},\\ldots,I\\\}iteratively usingCCmatrix multiplications of sizeD×DD\\times D\. Then, we evaluate the within\-chunk contribution via the batched contraction\. Finally, we advance the inter\-chunk state assin\(k\+1\)=AC​sin\(k\)\+contrib\(k\)s\_\{\\mathrm\{in\}\}^\{\(k\+1\)\}=A^\{C\}\\,s\_\{\\mathrm\{in\}\}^\{\(k\)\}\+\\text\{contrib\}^\{\(k\)\}\. SinceAAis frozen, the powers are computed on the fly and discarded after each chunk to avoid storing aC×D×DC\\times D\\times Dbuffer\.

ApproachSST\-2Poem SentimentAdditionCoin FlipEdge ExistenceAlgorithm[2](https://arxiv.org/html/2609.17888#alg2)95\.9±0\.795\.9\_\{\\pm 0\.7\}78\.6±1\.178\.6\_\{\\pm 1\.1\}81\.8±1\.281\.8\_\{\\pm 1\.2\}85\.3±1\.285\.3\_\{\\pm 1\.2\}83\.6±0\.683\.6\_\{\\pm 0\.6\}Algorithm[2](https://arxiv.org/html/2609.17888#alg2)without context compression95\.0±0\.595\.0\_\{\\pm 0\.5\}79\.7±0\.879\.7\_\{\\pm 0\.8\}82\.8±1\.082\.8\_\{\\pm 1\.0\}85\.3±1\.285\.3\_\{\\pm 1\.2\}83\.7±1\.383\.7\_\{\\pm 1\.3\}Table 8:Test accuracy \(%\) of Algorithm[2](https://arxiv.org/html/2609.17888#alg2)with and without context compression\. Each result reports the mean and standard deviation over three random seeds\.ApproachWindowMMLUModular AdditionGSM8KLengthErrorFLOPsMemoryErrorFLOPsMemoryErrorFLOPsMemoryStreamingLLM32322\.5%2\.5\\%7\.72​e137\.72e^\{13\}15\.8815\.88G10\.7%10\.7\\%1\.90​e131\.90e^\{13\}13\.9513\.95G6\.4%6\.4\\%1\.25​e141\.25e^\{14\}14\.0514\.05G1281280\.7%0\.7\\%7\.75​e137\.75e^\{13\}15\.9515\.95G8\.7%8\.7\\%1\.91​e131\.91e^\{13\}13\.9613\.96G5\.1%5\.1\\%1\.26​e141\.26e^\{14\}14\.0714\.07GLM\-Infinite323211\.6%11\.6\\%7\.73​e137\.73e^\{13\}15\.9015\.90G14\.3%14\.3\\%1\.90​e131\.90e^\{13\}12\.9612\.96G5\.7%5\.7\\%1\.26​e141\.26e^\{14\}14\.0614\.06G1281285\.1%5\.1\\%7\.75​e137\.75e^\{13\}15\.9815\.98G8\.7%8\.7\\%1\.91​e131\.91e^\{13\}12\.9912\.99G4\.9%4\.9\\%1\.27​e141\.27e^\{14\}14\.0714\.07G

Table 9:We report the results using more window lengths on MMLU, Addition datasets, and GSM8K withk=50k=50\. The relative error is computed against actual inference results\.This reduces the number of sequential state updates fromTTto⌈T/C⌉\\lceil T/C\\rceil, while each chunk\-level operation is a single, well\-parallelised GPU kernel\. The underlying SSM dynamics are unchanged—the final statesTs\_\{T\}is mathematically identical to that of the token\-by\-token scan\.

During training, we observe that directly optimizing the distillation objectives \(logits and hidden states\) without prior initialization degrades performance\. The root cause is a bootstrap problem: randomly initialized virtual KV tokens produce degenerate attention patterns in the frozen Transformer, which in turn yield near\-uniform output distributions\. The resulting logit\-level and hidden\-state gradients carry little useful signal for the auxiliary, preventing meaningful learning\. The first\-stage KV alignment resolves this by constraining the virtual tokens to lie in a geometrically valid region of the KV representation space before any query\-conditioned supervision is applied\.

### B\.3Omitted Experiment Results

We illustrate additional evaluation results across different model scales and sequence lengths to supplement the findings in the main text\.

#### Sensitivity to small numbers of demonstrations\.

To verify on shorter context lengths, Table[7](https://arxiv.org/html/2609.17888#A2.T7)extends the relative output error analysis to the context window ofk=25k=25demonstrations\. Across two different models \(Qwen\-1\.5B, Qwen\-3B\) and four downstream tasks, the relative error between the full demonstration prefix and our compressed virtual tokens remains below0\.6%0\.6\\%\. These results confirm that Algorithm[1](https://arxiv.org/html/2609.17888#alg1)consistently preserves the original predictive distribution regardless of the context scale or model architecture\.

#### Results on smaller models\.

To test whether Algorithm[1](https://arxiv.org/html/2609.17888#alg1)maintains its advantage at a smaller model scale, Table[10](https://arxiv.org/html/2609.17888#A2.T10)details the computational cost, memory usage, and relative error for the Qwen\-1\.5B\-Instruct model on the SST\-2, Poem Sentiment, and Modular Addition datasets\. Under the standard setting ofk=50k=50demonstrations, Algorithm[1](https://arxiv.org/html/2609.17888#alg1)achieves the lowest relative error among all evaluated methods, including0\.3%0\.3\\%on SST\-2 and0\.1%0\.1\\%on both Poem Sentiment and Addition\. Furthermore, Algorithm[1](https://arxiv.org/html/2609.17888#alg1)requires fewer FLOPs than the leading baseline, BSA, while maintaining a comparable memory footprint\. This verifies that the efficiency and accuracy advantages of the Algorithm[1](https://arxiv.org/html/2609.17888#alg1)module generalize reliably to smaller model variants\.

#### Coverage of remaining task categories\.

To confirm the extension to other reasoning tasks, Table[10](https://arxiv.org/html/2609.17888#A2.T10)and Table[11](https://arxiv.org/html/2609.17888#A2.T11)present the evaluation on the remaining reasoning datasets, Coin Flip and Edge Existence, using both the Qwen\-1\.5B\-Instruct and Qwen\-3B\-Instruct models withk=50k=50demonstrations\. Consistent with earlier findings, Algorithm[1](https://arxiv.org/html/2609.17888#alg1)achieves the lowest relative error among all baselines, recording0\.1%0\.1\\%error on both datasets for the 1\.5B model, and under0\.9%0\.9\\%for the 3B model\. This demonstrates that Algorithm[1](https://arxiv.org/html/2609.17888#alg1)preserves strong predictive fidelity across diverse reasoning tasks while maintaining a highly efficient computational profile compared to sliding\-window and block\-sparse baselines\.

#### Results regarding Algorithm[2](https://arxiv.org/html/2609.17888#alg2)\.

To separate the effect of the selection algorithm from that of context compression, in Table[8](https://arxiv.org/html/2609.17888#A2.T8), we find that removing context compression changes the average accuracy from85\.04%85\.04\\%to85\.30%85\.30\\%, corresponding to an increase of0\.260\.26percentage points\. This result shows that the accuracy gain cannot be explained by the compressed representation acting as a regularizer\. The selection algorithm remains effective when candidates are evaluated using the full context\.

MethodWindowSST\-2Poem SentimentModular AdditionLengthErrorFLOPsMemoryErrorFLOPsMemoryErrorFLOPsMemoryStreamingLLM323211\.3%11\.3\\%9\.42​e129\.42e^\{12\}6\.946\.94G17\.5%17\.5\\%1\.36​e131\.36e^\{13\}6\.616\.61G12\.8%12\.8\\%9\.53​e129\.53e^\{12\}7\.217\.21G64647\.6%7\.6\\%9\.43​e129\.43e^\{12\}6\.946\.94G13\.2%13\.2\\%1\.36​e131\.36e^\{13\}6\.626\.62G11\.9%11\.9\\%9\.54​e129\.54e^\{12\}7\.227\.22G1281285\.2%5\.2\\%9\.46​e129\.46e^\{12\}6\.966\.96G8\.7%8\.7\\%1\.37​e131\.37e^\{13\}6\.646\.64G8\.4%8\.4\\%9\.56​e129\.56e^\{12\}7\.247\.24GLM\-Infinite323216\.2%16\.2\\%9\.36​e129\.36e^\{12\}6\.936\.93G7\.0%7\.0\\%1\.37​e131\.37e^\{13\}6\.616\.61G5\.5%5\.5\\%9\.41​e129\.41e^\{12\}7\.217\.21G646411\.5%11\.5\\%9\.38​e129\.38e^\{12\}6\.946\.94G3\.7%3\.7\\%1\.37​e131\.37e^\{13\}6\.616\.61G3\.5%3\.5\\%9\.44​e129\.44e^\{12\}7\.217\.21G12812811\.0%11\.0\\%9\.41​e129\.41e^\{12\}6\.966\.96G2\.9%2\.9\\%1\.37​e131\.37e^\{13\}6\.646\.64G3\.2%3\.2\\%9\.50​e129\.50e^\{12\}7\.237\.23GBSA∼400\\sim 4001\.1%1\.1\\%9\.52​e129\.52e^\{12\}7\.157\.15G0\.7%0\.7\\%1\.53​e131\.53e^\{13\}6\.966\.96G0\.9%0\.9\\%8\.81​e128\.81e^\{12\}7\.787\.78GAlgorithm[1](https://arxiv.org/html/2609.17888#alg1)𝟏𝟔\\mathbf\{16\}0\.3%\\mathbf\{0\.3\\%\}6\.30​𝐞𝟏𝟐\\mathbf\{6\.30e^\{12\}\}7\.19\\mathbf\{7\.19\}G0\.1%\\mathbf\{0\.1\\%\}1\.02​𝐞𝟏𝟑\\mathbf\{1\.02e^\{13\}\}7\.06\\mathbf\{7\.06\}G0\.1%\\mathbf\{0\.1\\%\}5\.28​𝐞𝟏𝟐\\mathbf\{5\.28e^\{12\}\}8\.06\\mathbf\{8\.06\}GMethodWindowCoin FlipEdge ExistenceMMLULengthErrorFLOPsMemoryErrorFLOPsMemoryErrorFLOPsMemoryStreamingLLM323218\.1%18\.1\\%1\.30​e131\.30e^\{13\}7\.837\.83G3\.8%3\.8\\%5\.99​e135\.99e^\{13\}14\.6314\.63G3\.2%3\.2\\%3\.88​e133\.88e^\{13\}10\.0610\.06G646410\.2%10\.2\\%1\.30​e131\.30e^\{13\}7\.847\.84G2\.7%2\.7\\%6\.01​e136\.01e^\{13\}14\.6414\.64G0\.8%0\.8\\%3\.88​e133\.88e^\{13\}10\.0810\.08G1281289\.8%9\.8\\%1\.31​e131\.31e^\{13\}7\.857\.85G2\.5%2\.5\\%6\.03​e136\.03e^\{13\}14\.6514\.65G0\.6%0\.6\\%3\.89​e133\.89e^\{13\}10\.1310\.13GLM\-Infinite323228\.7%28\.7\\%1\.32​e131\.32e^\{13\}7\.847\.84G7\.5%7\.5\\%6\.43​e136\.43e^\{13\}14\.6314\.63G1\.4%1\.4\\%3\.88​e133\.88e^\{13\}10\.0710\.07G646426\.8%26\.8\\%1\.32​e131\.32e^\{13\}7\.847\.84G6\.5%6\.5\\%6\.43​e136\.43e^\{13\}14\.6414\.64G0\.7%0\.7\\%3\.89​e133\.89e^\{13\}10\.0810\.08G12812823\.9%23\.9\\%1\.33​e131\.33e^\{13\}7\.867\.86G4\.6%4\.6\\%6\.46​e136\.46e^\{13\}14\.6614\.66G0\.5%0\.5\\%3\.89​e133\.89e^\{13\}10\.1310\.13GBSA∼400\\sim 4000\.7%0\.7\\%1\.35​e131\.35e^\{13\}8\.288\.28G0\.9%0\.9\\%6\.88​e136\.88e^\{13\}15\.7715\.77G0\.5%0\.5\\%3\.95​e133\.95e^\{13\}10\.3610\.36GAlgorithm[1](https://arxiv.org/html/2609.17888#alg1)𝟏𝟔\\mathbf\{16\}0\.1%\\mathbf\{0\.1\\%\}9\.02​𝐞𝟏𝟐\\mathbf\{9\.02e^\{12\}\}8\.35\\mathbf\{8\.35\}G0\.1%\\mathbf\{0\.1\\%\}4\.11​𝐞𝟏𝟑\\mathbf\{4\.11e^\{13\}\}14\.73\\mathbf\{14\.73\}G0\.3%\\mathbf\{0\.3\\%\}2\.49​𝐞𝟏𝟑\\mathbf\{2\.49e^\{13\}\}10\.81\\mathbf\{10\.81\}G

Table 10:Supplementary results on Qwen\-1\.5B\-Instruct\. We compare long\-context inference methods on SST\-2, Poem Sentiment, Modular Addition, Coin Flip, Edge Existence, and MMLU withk=50k=50demonstrations, and report relative output error with respect to full\-prefix inference, together with FLOPs and peak memory usage\.MethodWindowCoin FlipEdge ExistencePoem SentimentLengthErrorFLOPsMemoryErrorFLOPsMemoryErrorFLOPsMemoryDense\-\-1\.49​e151\.49e^\{15\}14\.5814\.58G\-9\.77​e159\.77e^\{15\}24\.5724\.57G\-2\.53​e152\.53e^\{15\}12\.7012\.70GStreamingLLM646431\.5%31\.5\\%2\.63​e132\.63e^\{13\}13\.5913\.59G1\.9%1\.9\\%1\.28​e141\.28e^\{14\}20\.4020\.40G9\.3%9\.3\\%2\.72​e132\.72e^\{13\}12\.3712\.37G12812830\.6%30\.6\\%2\.64​e132\.64e^\{13\}13\.6113\.61G1\.8%1\.8\\%1\.28​e141\.28e^\{14\}20\.4220\.42G9\.1%9\.1\\%2\.73​e132\.73e^\{13\}12\.3912\.39GLM\-Infinite646432\.8%32\.8\\%2\.64​e132\.64e^\{13\}13\.6013\.60G1\.3%1\.3\\%1\.28​e141\.28e^\{14\}20\.4020\.40G14\.2%14\.2\\%2\.72​e132\.72e^\{13\}12\.3812\.38G12812830\.6%30\.6\\%2\.64​e132\.64e^\{13\}13\.6113\.61G1\.2%1\.2\\%1\.28​e141\.28e^\{14\}20\.4320\.43G13\.8%13\.8\\%2\.73​e132\.73e^\{13\}12\.4012\.40GICAE1281281\.4%1\.4\\%1\.36​e141\.36e^\{14\}14\.2714\.27G3\.5%3\.5\\%1\.02​e151\.02e^\{15\}24\.7624\.76G3\.2%3\.2\\%5\.40​e135\.40e^\{13\}12\.3212\.32GDuoAttention646410\.5%10\.5\\%2\.70​e132\.70e^\{13\}14\.0214\.02G3\.2%3\.2\\%1\.40​e141\.40e^\{14\}22\.5922\.59G7\.6%7\.6\\%2\.75​e132\.75e^\{13\}12\.5012\.50GBSA∼400\\sim 4002\.4%2\.4\\%2\.75​e132\.75e^\{13\}14\.0714\.07G1\.4%1\.4\\%1\.53​e141\.53e^\{14\}24\.6424\.64G4\.2%4\.2\\%2\.95​e132\.95e^\{13\}12\.8512\.85GAlgorithm[1](https://arxiv.org/html/2609.17888#alg1)𝟏𝟔\\mathbf\{16\}0\.7%\\mathbf\{0\.7\\%\}1\.91​𝐞𝟏𝟑\\mathbf\{1\.91e^\{13\}\}14\.13\\mathbf\{14\.13\}G0\.8%\\mathbf\{0\.8\\%\}9\.97​𝐞𝟏𝟑\\mathbf\{9\.97e^\{13\}\}25\.02\\mathbf\{25\.02\}G1\.6%\\mathbf\{1\.6\\%\}2\.37​𝐞𝟏𝟑\\mathbf\{2\.37e^\{13\}\}13\.23\\mathbf\{13\.23\}G

Table 11:Supplementary results on Qwen\-3B\-Instruct\. We compare long\-context inference methods on the Coin Flip, Edge Existence, and Poem Sentiment tasks withk=50k=50demonstrations, and report relative output error with respect to full\-prefix inference, together with FLOPs and peak memory usage\.