cs.LGOct 14, 2026

Provable Transformer Context Compression for In-Context Reinforcement Learning

Authors: Yajie Zhao, Yingjun Mo, Zhiyong Wang, Bo Tang, Chao Xu, Fang Kong

Organizations: Southern University of Science and Technology · University of Electronic Science and Technology of China · Harbin Institute of Technology

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Transformers can perform in-context reinforcement learning without parameter updates. However, existing theoretical constructions often require context sizes that grow with the number of sampled transitions, leading to increasing memory costs in large-scale problems. This raises a fundamental scalability question: how much context is intrinsically necessary to carry out such learning? We answer this question for fixed-policy TD evaluation. We show that, with a shared linear decoder, the minimum context dimension required to exactly reproduce the first LL TD iterates is the Krylov dimension induced by the sampled transitions. When the available context dimension mm falls below this threshold, exact reproduction is no longer possible. In this regime, we characterize the optimal approximation error and show that it decays as Θ(γmmax⁡{1,1/(m(1−γ))})Θ(γ^m\max\{1,1/(m(1-γ))\}) where γγ is the discounting factor. Moving from representation complexity to actual memory cost, we characterize the optimal retained memory for exact continuous representations as Θ(nkL)Θ(nk_L) and show experimentally that on a sampled task whose exact rank is full but whose numerical rank is small, a context accurate to 10−210^{-2} is 88.7%88.7\% smaller than the sparse sufficient statistics of TD. Finally, we establish the learnability of the shared executor through pretraining by proving global convergence of an auxiliary objective over effective attention parameters under suitable task diversity. Experiments across diverse TD evaluation problems validate parameter recovery and the predicted context--accuracy tradeoff.

Figures & tables

WorkProblem settingRetained objectRetained memoryBuffer (KiB)
Wang et al. (2025a)TD(0), linear features, sampled trajectoryNN transition tokens (ϕt,γϕt+1,Rt+1)(\phi_{t},\gamma\phi_{t+1},R_{t+1})Θ(Nd)\Theta(Nd)383.00
Wang et al. (2025b)tabular TD, known dynamicsnn state tokens with rows of γP\gamma P and rrΘ(n2)\Theta(n^{2})4.79
Xie et al. (2026b)softmax (kernel) TD, sampled trajectoryNN state tokens (xt,Rt+1)(x_{t},R_{t+1})Θ(Nd)\Theta(Nd)211.04
Xie et al. (2026a)TD(0) with chain of thought, linear featuresNN transition tokens and LL iterate tokensΘ((N+L)d)\Theta((N+L)d)579.15
Ourstabular TD (Thm. 4.3 ) or linear features (Thm. 4.4 ), sampled trajectorykLk_{L} Krylov tokens (B,Y,g)(B,Y,g)O(nkL+n)\mathbf{O(nk_{L}+n)}2.49\mathbf{2.49}
Table 1: Constructions that execute TD policy evaluation in the forward pass. N,n,d,LN,n,d,L denote transitions, states, feature dimension, and TD depth; kL=dim⁡KL(γP^,r^)k_{L}=\dim\mathcal{K}_{L}(\gamma\widehat{P},\widehat{r}) . Theory counts task scalars; buffers are float32 array bytes /1024/1024 over 100 encodings at N=2000N=2000 , n=d=24n=d=24 , L=30L=30 , kL=12k_{L}=12 , excluding parameters and forward scratch. Constructions for other in-context RL problems ( Lin et al., 2024 ; Chen et al., 2025 ; Liang and Lai, 2026 ) retain different objects and are listed in Appendix G as not comparable.
Figure 1: Parameter identification. Learned C,WC,W : calibration ( 7 ) with full-batch gradient descent (a,b), trajectory supervision with Adam (c,d), and the depth-30 endpoint surrogate with Adam (e,f).
Figure 2: Context-dimension–accuracy tradeoff. (a) Cyclic MRPs with kL=12k_{L}=12 and (b) reversible MRPs: mean error over 100 tasks per panel ( n=24n=24 , γ=0.9\gamma=0.9 , L=30L=30 ), interquartile ranges shaded, dashed lines at kLk_{L} . (c) Worst-case rings at γ=0.99\gamma=0.99 against Ψm\Psi_{m} of Theorem 4.5 and a certified lower bound (vertical line: m=1/(1−γ)m=1/(1-\gamma) ).
Tasknn or ddNNRawStatisticsSparse TDOursPeak
Block cycle2496011.442.440.752.4937.5
Block cycle48019200228.8903.8102.248.810822.5
Boyan’s chain420000703.10.08–0.181.1
Mountain car45232000021250799.8–116.79226
Table 2: Retained memory (KiB, float32). Top: 12-block cycles with kL=12k_{L}=12 . Bottom: feature TD ( 6 ), Statistics =(AN,bN)=(A_{N},b_{N}) , Ours == the exact context (Boyan) or the dimension- 3232 context with error ≤1.1×10−2\leq 1.1\times 10^{-2} (mountain car; Raw stored sparsely). Peak as in Definition 4.2 .
Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

CheckMax
Exact Krylov context, softmax executor with Λn=2n(n+1)\Lambda_{n}=2n(n+1) vs. direct TD5.3×10−135.3\times 10^{-13}
same vs. the linear-attention executor on the same context5.3×10−135.3\times 10^{-13}
Budget- mm Galerkin context, softmax vs. linear executor8.9×10−138.9\times 10^{-13}
Sample-and-hold, softmax clocked head with Λm,γ=2m/(1−γ)\Lambda_{m,\gamma}=2m/(1-\gamma) vs. v^σ(ℓ)\widehat{v}_{\sigma(\ell)}6.0×10−146.0\times 10^{-14}
Krylov-anchored hold, softmax vs. linear executor (family gain from the context)1.5×10−141.5\times 10^{-14}
Largest gain ratio, Krylov / clocked0.110.11 / 0.250.25
Appendix
Table 3: Softmax realization checks on 60 random covered batches ( 4≤n≤244\leq n\leq 24 , γ∈{0.5,0.8,0.9,0.95}\gamma\in\{0.5,0.8,0.9,0.95\} , 10≤L≤6010\leq L\leq 60 , ∥r^∥∞=1\|\widehat{r}\|_{\infty}=1 ). Errors are maximum absolute deviations divided by max⁡ℓ,s∣v^ℓ(s)∣\max_{\ell,s}|\widehat{v}_{\ell}(s)| . The gain ratio is the family quantity in ( 11 ) divided by the task-uniform gain of Corollary B.2 ; all 60 batches were admissible.
Check (instances)Max
Stage-1 heads vs. direct counting of (A^,r^)(\widehat{A},\widehat{r}) (60 random batches)2.2×10−162.2\times 10^{-16}
Exact Krylov context + TD executor vs. direct TD, balanced two-pass (60)3.8×10−153.8\times 10^{-15} rel.
same, literal unbalanced one-pass encoder (60)8.0×10−138.0\times 10^{-13} rel.
Visit-weighted Galerkin execution vs. closed form; dual biorthogonality (60)2.4×10−152.4\times 10^{-15} ; 4.4×10−164.4\times 10^{-16}
Sample-and-hold encoder + depth-counter executor vs. v^σ(ℓ)\widehat{v}_{\sigma(\ell)} (60)8.9×10−168.9\times 10^{-16}
Krylov-anchored encoder + executor vs. v~ℓG+eσ(ℓ)\widetilde{v}^{\rm G}_{\ell}+e_{\sigma(\ell)} (40)2.8×10−152.8\times 10^{-15}
Appendix
Table 4: Implementation identities and bound checks. “Max” is the maximum over the stated random instances. Ratios below one mean that the bound holds with room to spare.
Figure 3: γ\gamma -uniform context–accuracy law. Worst normalized all-depth error over structured ring instances, built from raw batches through attention heads, at γ=0.9\gamma=0.9 and γ=0.99\gamma=0.99 . Triangles give a certified lower bound for any mm -dimensional context; the vertical line marks m=1/(1−γ)m=1/(1-\gamma) . Sample-and-hold and the anchored hold follow Ψm\Psi_{m} , whereas Krylov truncation plateaus until the effective horizon. The right panel is reproduced in Figure 2 (c).
Figure 4: Sampled noisy tasks. Mean normalized error over five raw trajectories versus the context dimension. The horizontal line is the sampling error of the empirical trajectory v^\widehat{v} relative to the population TD trajectory; the band is the numerical-rank window of Proposition D.3 (b).
TasknnRawStatisticsSparse TDAnswersOursLB
Block cycle, exact2411.442.440.759.382.491.17
Block cycle, exact240114.4226.936.293.824.411.3
Block cycle, exact480228.8903.8102.2187.548.822.5
Dense kernel, ϵ=10−2\epsilon=10^{-2}2002345157.8139.0234.415.7 / 17.2–
Goal gridworld, ϵ=10−2\epsilon=10^{-2}126117363.05.7147.738.5 / 44.5–
Appendix
Table 5: Retained memory (KiB, float32) on exact and sampled tasks. Exact rows: dense-block cycles, γ=0.9\gamma=0.9 , L=100L=100 , N=40nN=40n ; LB is nkL+kLnk_{L}+k_{L} (Theorem 4.3 (c)). Approximate rows: γ=0.99\gamma=0.99 , L=300L=300 , smallest context dimension with error ≤10−2\leq 10^{-2} (Krylov projection / anchored hold). Bold marks the smallest entry in each row.
nnRetainedStatisticsPeak, Stage 1of which chunkPeak, Stage 2Peak, executionTotal
242.492.4428.19.437.53.137.5
969.8036.75436.5145.5582.010.4582.0
24024.43226.882711.3903.81947.225.02711.3
48048.80903.7510822.53607.55985.049.410822.5
Appendix
Table 6: Retained and peak live storage (KiB, float32) of the exact dense-block runs, with Stage 1 streamed in chunks of 2n2n records; “chunk” is the part of the Stage-1 peak occupied by one chunk of one-hot records, and both stage peaks include the heads’ working storage. Execution is the analytic budget of one query. Total is the peak task memory of Definition 4.2 .
 RetainedTotal peakMax normalized error
nnfloat32float64float32float64float64float32 contextfloat32 pipeline
242.494.9837.575.01.6×10−161.6\times 10^{-16}2.0×10−72.0\times 10^{-7}2.0×10−72.0\times 10^{-7}
969.8019.60582.01164.01.7×10−151.7\times 10^{-15}2.4×10−72.4\times 10^{-7}2.7×10−72.7\times 10^{-7}
24024.4348.852711.35422.59.6×10−169.6\times 10^{-16}6.2×10−76.2\times 10^{-7}4.5×10−74.5\times 10^{-7}
48048.8097.6010822.521645.02.9×10−152.9\times 10^{-15}3.3×10−73.3\times 10^{-7}4.2×10−74.2\times 10^{-7}
Appendix
Table 7: Memory and TD error under one data type. Retained memory of the Krylov context and total peak in float32 and float64 (KiB), and the maximum absolute error of the executed trajectory relative to max⁡ℓ,s∣v^ℓ(s)∣\max_{\ell,s}|\widehat{v}_{\ell}(s)| for the float64 pipeline, for the float64 context stored and executed in float32, and for the whole pipeline in float32.
ε\varepsilonRank L/BPrompt LPrompt BError LError B
001/11.01.02.02.03.6×10−153.6\times 10^{-15}00
10−410^{-4}2/21.1×1041.1\times 10^{4}2.02.02.0×10−122.0\times 10^{-12}3.6×10−153.6\times 10^{-15}
10−810^{-8}2/21.1×1081.1\times 10^{8}2.02.01.5×10−81.5\times 10^{-8}8.9×10−168.9\times 10^{-16}
10−1210^{-12}2/11.1×10121.1\times 10^{12}2.02.01.5×10−121.5\times 10^{-12}3.1×10−123.1\times 10^{-12}
10−1610^{-16}1/11.01.02.02.03.6×10−153.6\times 10^{-15}00
Appendix
Table 8: Rank-boundary sensitivity. Retained rank (L/B), largest context entry, and maximum all-state, all-depth error. At ε=10−16\varepsilon=10^{-16} the float64 input already equals one.
Figure 5: Learned bottleneck encoders against compiled contexts. (a) Mean normalized error on 256 held-out tasks and (b) worst error over three ring instances, for standard Transformer encoders with a clocked or time-invariant executor (three seeds per point, 8,000 updates each; error bars are one standard deviation across seeds), the compiled contexts on the same batches, Ψm\Psi_{m} , and the certified lower bound for every mm -dimensional context. The vertical line marks m=1/(1−γ)m=1/(1-\gamma) .
 Held-out tasks, mean errorRing instances, worst error
mmclockedtime-inv.SHanchoredclockedtime-inv.SHe2Ψme^{2}\Psi_{m}LB
10.89±0.010.89\pm 0.01–1.721.722.172.171.64±0.011.64\pm 0.01–3.083.0866.5066.500.530.53
20.69±0.000.69\pm 0.000.64±0.110.64\pm 0.111.391.391.851.851.30±0.011.30\pm 0.011.26±0.021.26\pm 0.022.282.2829.9329.930.310.31
40.61±0.030.61\pm 0.030.56±0.040.56\pm 0.041.231.230.930.931.16±0.031.16\pm 0.031.12±0.111.12\pm 0.111.731.7312.1212.120.230.23
80.63±0.030.63\pm 0.030.51±0.000.51\pm 0.001.001.000.330.331.14±0.021.14\pm 0.020.98±0.040.98\pm 0.041.001.003.983.980.140.14
160.62±0.060.62\pm 0.060.50±0.000.50\pm 0.000.230.230.0870.0871.15±0.061.15\pm 0.060.95±0.050.95\pm 0.050.430.431.371.370.0600.060
Appendix
Table 9: Learned bottleneck encoders (trajectory-token interface; mean ±\pm one standard deviation over three seeds, 8,0008{,}000 updates each) against the compiled sample-and-hold (SH) and anchored contexts on the same batches, the upper bound e2Ψme^{2}\Psi_{m} of Theorem 4.5 , and the certified lower bound (LB) of the ring instances. Dashes: budget not trained, or no certificate.
More in this topic

In-Context RL

In-context reinforcement learning studies how agents adapt their policies from interaction histories without updating model parameters.

See all 42 papers

Explore similar work

CardsList
May 7, 2026

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

We investigate the ability of transformers to perform in-context reinforcement learning (ICRL), where a model must infer and execute learning algorithms from trajectory data without parameter updates. We show that a linear self-attention transformer block can provably implement policy-improvement methods, including semi-gradient SARSA and actor-critic, via explicit parameter constructions. Beyond existence, we design a teacher-mimicking training procedure, analyze its gradient-flow dynamics, and establish the first convergence guarantee in the ICRL literature: under suitable richness conditions on the training MDP distribution, gradient flow converges locally and exponentially to an optimal parameter manifold corresponding to the desired RL update. Empirically, training transformers on randomly generated tabular MDPs confirms these predictions: the learned models recover the parameter structure of our explicit constructions and, when deployed on unseen MDPs, deliver strong in-context control performance. Together, these results illuminate how transformer architectures internalize and execute classical reinforcement learning algorithms in context, bridging mechanistic understanding and training dynamics in ICRL.

Haodong Liang, Lifeng Lai

Department of Electrical and Computer Engineering, University of California, Davis

Reinforcement LearningTransformer-Based RLPolicy OptimizationIn-Context RLPolicy Learning
AskPDF
May 8, 2026

Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning

In-context reinforcement learning (ICRL) studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates. Existing theoretical analyses of ICRL largely rely on linear attention, which replaces the softmax function in the standard attention with an identity mapping. This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification. In particular, we consider the standard softmax attention used in practice. We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm. Here, weighted softmax TD is a new RL algorithm that performs policy evaluation in kernel space and adopts both linear TD and tabular TD as special cases. We also prove that under a certain contraction condition, the policy evaluation error decays as the number of layers grows, with the identified parameters above. Finally, we prove that those parameters are a global minimizer of a pretraining loss, explaining their emergence in our numerical experiments.

Zixuan Xie, Xinyu Liu, Claire Chen +3

University of Virginia · California Institute of Technology · Purdue University

Reinforcement LearningTransformer-Based RLTemporal-Difference LearningTransformer AttentionIn-Context RLContextual RL
AskPDF
May 8, 2026

Convergence and Emergence of In-Context Reinforcement Learning with Chain of Thought

In-context reinforcement learning (ICRL) refers to the ability of RL agents to adapt to new tasks at inference time without parameter updates by conditioning on additional context. Recent empirical studies further demonstrate that Chain-of-Thought (CoT) generation can amplify this ICRL capability. This paper is the first to provide a theoretical understanding on how CoT interacts with ICRL. We conduct our analysis in a policy evaluation setup with linear Transformer. We prove that with specific Transformer parameters, the CoT generation process is equivalent to repeatedly executing temporal difference learning updates. Additionally, we provide finite sample convergence analysis showing that the policy evaluation error decreases geometrically with CoT length and eventually saturates at a statistical floor determined by the context length. We also prove that the desired Transformer parameters are a global minimizer of the pretraining loss, providing a theoretical understanding on the empirical emergence of those parameters.

Zixuan Xie, Xinyu Liu, Rohan Chandra +1

University of Virginia

Reinforcement LearningTransformerTransformer-Based RLCoT ReasoningTemporal-Difference LearningIn-Context Learning+1
AskPDF