Transformers reduce memory needs for reinforcement learning context data
Provable Transformer Context Compression for In-Context Reinforcement Learning
Machine LearningArtificial Intelligence
Summary
Transformers can learn from examples they see without changing their brains, but usually need more and more past examples to remember as they get more data. The authors studied how much memory these models really need to learn a fixed strategy in reinforcement learning. They found a precise measure for the minimal memory size required and explained how errors get smaller as memory increases. Their results also showed that training transformers with shared components is possible and tested this theory successfully on various tasks.
What this means in practice
- •For machine learning engineers: Design models that minimize memory use for in-context reinforcement learning tasks by exploiting context compression limits characterized in the paper.
- •For ai system architects: Build scalable reinforcement learning systems that balance context size and accuracy based on the paper’s tradeoff between memory and approximation error.
Authors
Yajie Zhao, Yingjun Mo, Zhiyong Wang, Bo Tang, Chao Xu, Fang Kong
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 $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.