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
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 L TD iterates is the Krylov dimension induced by the sampled transitions. When the available context dimension m 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−γ))}) 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) and show experimentally that on a sampled task whose exact rank is full but whose numerical rank is small, a context accurate to 10−2 is 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
Work
Problem setting
Retained object
Retained memory
Buffer (KiB)
Wang et al. (2025a)
TD(0), linear features, sampled trajectory
N transition tokens (ϕt,γϕt+1,Rt+1)
Θ(Nd)
383.00
Wang et al. (2025b)
tabular TD, known dynamics
n state tokens with rows of γP and r
Θ(n2)
4.79
Xie et al. (2026b)
softmax (kernel) TD, sampled trajectory
N state tokens (xt,Rt+1)
Θ(Nd)
211.04
Xie et al. (2026a)
TD(0) with chain of thought, linear features
N transition tokens and L iterate tokens
Θ((N+L)d)
579.15
Ours
tabular TD (Thm. 4.3 ) or linear features (Thm. 4.4 ), sampled trajectory
kL Krylov tokens (B,Y,g)
O(nkL+n)
2.49
Table 1: Constructions that execute TD policy evaluation in the forward pass. N,n,d,L denote transitions, states, feature dimension, and TD depth; kL=dimKL(γP,r) . Theory counts task scalars; buffers are float32 array bytes /1024 over 100 encodings at N=2000 , n=d=24 , L=30 , kL=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,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=12 and (b) reversible MRPs: mean error over 100 tasks per panel ( n=24 , γ=0.9 , L=30 ), interquartile ranges shaded, dashed lines at kL . (c) Worst-case rings at γ=0.99 against Ψm of Theorem 4.5 and a certified lower bound (vertical line: m=1/(1−γ) ).
Task
n or d
N
Raw
Statistics
Sparse TD
Ours
Peak
Block cycle
24
960
11.44
2.44
0.75
2.49
37.5
Block cycle
480
19200
228.8
903.8
102.2
48.8
10822.5
Boyan’s chain
4
20000
703.1
0.08
–
0.18
1.1
Mountain car
452
320000
21250
799.8
–
116.7
9226
Table 2: Retained memory (KiB, float32). Top: 12-block cycles with kL=12 . Bottom: feature TD ( 6 ), Statistics =(AN,bN) , Ours = the exact context (Boyan) or the dimension- 32 context with error ≤1.1×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
Check
Max
Exact Krylov context, softmax executor with Λn=2n(n+1) vs. direct TD
5.3×10−13
same vs. the linear-attention executor on the same context
5.3×10−13
Budget- m Galerkin context, softmax vs. linear executor
8.9×10−13
Sample-and-hold, softmax clocked head with Λm,γ=2m/(1−γ) vs. vσ(ℓ)
6.0×10−14
Krylov-anchored hold, softmax vs. linear executor (family gain from the context)
1.5×10−14
Largest gain ratio, Krylov / clocked
0.11 / 0.25
Appendix
Table 3: Softmax realization checks on 60 random covered batches ( 4≤n≤24 , γ∈{0.5,0.8,0.9,0.95} , 10≤L≤60 , ∥r∥∞=1 ). Errors are maximum absolute deviations divided by maxℓ,s∣vℓ(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) (60 random batches)
2.2×10−16
Exact Krylov context + TD executor vs. direct TD, balanced two-pass (60)
3.8×10−15 rel.
same, literal unbalanced one-pass encoder (60)
8.0×10−13 rel.
Visit-weighted Galerkin execution vs. closed form; dual biorthogonality (60)
2.4×10−15 ; 4.4×10−16
Sample-and-hold encoder + depth-counter executor vs. vσ(ℓ) (60)
8.9×10−16
Krylov-anchored encoder + executor vs. vℓG+eσ(ℓ) (40)
2.8×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: γ -uniform context–accuracy law. Worst normalized all-depth error over structured ring instances, built from raw batches through attention heads, at γ=0.9 and γ=0.99 . Triangles give a certified lower bound for any m -dimensional context; the vertical line marks m=1/(1−γ) . Sample-and-hold and the anchored hold follow Ψ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 relative to the population TD trajectory; the band is the numerical-rank window of Proposition D.3 (b).
Task
n
Raw
Statistics
Sparse TD
Answers
Ours
LB
Block cycle, exact
24
11.44
2.44
0.75
9.38
2.49
1.17
Block cycle, exact
240
114.4
226.9
36.2
93.8
24.4
11.3
Block cycle, exact
480
228.8
903.8
102.2
187.5
48.8
22.5
Dense kernel, ϵ=10−2
200
2345
157.8
139.0
234.4
15.7 / 17.2
–
Goal gridworld, ϵ=10−2
126
1173
63.0
5.7
147.7
38.5 / 44.5
–
Appendix
Table 5: Retained memory (KiB, float32) on exact and sampled tasks. Exact rows: dense-block cycles, γ=0.9 , L=100 , N=40n ; LB is nkL+kL (Theorem 4.3 (c)). Approximate rows: γ=0.99 , L=300 , smallest context dimension with error ≤10−2 (Krylov projection / anchored hold). Bold marks the smallest entry in each row.
n
Retained
Statistics
Peak, Stage 1
of which chunk
Peak, Stage 2
Peak, execution
Total
24
2.49
2.44
28.1
9.4
37.5
3.1
37.5
96
9.80
36.75
436.5
145.5
582.0
10.4
582.0
240
24.43
226.88
2711.3
903.8
1947.2
25.0
2711.3
480
48.80
903.75
10822.5
3607.5
5985.0
49.4
10822.5
Appendix
Table 6: Retained and peak live storage (KiB, float32) of the exact dense-block runs, with Stage 1 streamed in chunks of 2n 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 .
Retained
Total peak
Max normalized error
n
float32
float64
float32
float64
float64
float32 context
float32 pipeline
24
2.49
4.98
37.5
75.0
1.6×10−16
2.0×10−7
2.0×10−7
96
9.80
19.60
582.0
1164.0
1.7×10−15
2.4×10−7
2.7×10−7
240
24.43
48.85
2711.3
5422.5
9.6×10−16
6.2×10−7
4.5×10−7
480
48.80
97.60
10822.5
21645.0
2.9×10−15
3.3×10−7
4.2×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)∣ for the float64 pipeline, for the float64 context stored and executed in float32, and for the whole pipeline in float32.
ε
Rank L/B
Prompt L
Prompt B
Error L
Error B
0
1/1
1.0
2.0
3.6×10−15
0
10−4
2/2
1.1×104
2.0
2.0×10−12
3.6×10−15
10−8
2/2
1.1×108
2.0
1.5×10−8
8.9×10−16
10−12
2/1
1.1×1012
2.0
1.5×10−12
3.1×10−12
10−16
1/1
1.0
2.0
3.6×10−15
0
Appendix
Table 8: Rank-boundary sensitivity. Retained rank (L/B), largest context entry, and maximum all-state, all-depth error. At ε=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 , and the certified lower bound for every m -dimensional context. The vertical line marks m=1/(1−γ) .
Held-out tasks, mean error
Ring instances, worst error
m
clocked
time-inv.
SH
anchored
clocked
time-inv.
SH
e2Ψm
LB
1
0.89±0.01
–
1.72
2.17
1.64±0.01
–
3.08
66.50
0.53
2
0.69±0.00
0.64±0.11
1.39
1.85
1.30±0.01
1.26±0.02
2.28
29.93
0.31
4
0.61±0.03
0.56±0.04
1.23
0.93
1.16±0.03
1.12±0.11
1.73
12.12
0.23
8
0.63±0.03
0.51±0.00
1.00
0.33
1.14±0.02
0.98±0.04
1.00
3.98
0.14
16
0.62±0.06
0.50±0.00
0.23
0.087
1.15±0.06
0.95±0.05
0.43
1.37
0.060
Appendix
Table 9: Learned bottleneck encoders (trajectory-token interface; mean ± one standard deviation over three seeds, 8,000 updates each) against the compiled sample-and-hold (SH) and anchored contexts on the same batches, the upper bound e2Ψm of Theorem 4.5 , and the certified lower bound (LB) of the ring instances. Dashes: budget not trained, or no certificate.