Provable Transformer Context Compression for In-Context Reinforcement Learning

📄 arXiv: 2610.20004v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-10-14

💡 一句话要点

提出可证明的Transformer上下文压缩以解决强化学习中的内存问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 上下文压缩 强化学习 时间差分 Transformer 内存优化 Krylov维度 参数恢复

📋 核心要点

  1. 现有方法在进行上下文强化学习时，往往需要随着采样过渡数量的增加而增加上下文大小，导致内存成本高昂。
  2. 论文提出了一种基于共享线性解码器的框架，确定了进行TD评估所需的最小上下文维度，并分析了在上下文维度不足时的近似误差。
  3. 实验结果表明，在某些任务中，准确到10^{-2}的上下文比TD的稀疏充分统计量小88.7%，验证了所提出方法的有效性。

📝 摘要（中文）

论文探讨了在不更新参数的情况下，Transformer如何进行上下文强化学习。现有理论构建通常要求上下文大小随着采样过渡的数量增长，从而导致大规模问题中的内存成本增加。本文针对固定策略的时间差分（TD）评估，确定了进行学习所需的最小上下文维度，并通过实验验证了参数恢复和上下文-准确性之间的权衡。

🔬 方法详解

问题定义：本文解决了在固定策略TD评估中，如何有效压缩上下文以降低内存成本的问题。现有方法在大规模问题中面临上下文大小与内存成本之间的矛盾。

核心思路：论文通过共享线性解码器，确定了重现TD迭代所需的最小上下文维度，并在上下文不足时分析了近似误差的衰减特性。

技术框架：整体框架包括上下文维度的计算、近似误差的分析以及通过预训练实现共享执行器的学习。主要模块包括上下文生成、误差评估和参数恢复。

关键创新：最重要的创新在于通过Krylov维度的引入，明确了上下文维度与TD迭代重现之间的关系，提供了理论支持。与现有方法相比，提供了更为精确的上下文需求分析。

关键设计：关键设计包括上下文维度的选择、损失函数的设定以及有效注意力参数的优化，确保了在不同任务下的全局收敛性。实验中使用的参数设置和网络结构经过精心设计，以适应多样化的任务需求。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

实验结果显示，所提出的方法在某些任务中能够将上下文的内存需求减少88.7%，同时保持准确度在10^{-2}的范围内，验证了上下文与准确性之间的有效权衡。

🎯 应用场景

该研究的潜在应用领域包括强化学习、自动控制和机器人技术等。通过优化上下文使用，能够在内存受限的环境中实现更高效的学习，未来可能推动大规模强化学习系统的实际应用。

📄 摘要（原文）

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 $Θ(γ^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 $Θ(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^{-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.