arXiv:2506.19686cs.AI2025-06被引 1

模型通过记忆中间计算实现快速决策,模拟生物学习机制。

From Memories to Maps: Mechanisms of In-Context Reinforcement Learning in Transformers

  • 用记忆令牌缓存中间计算,决策时调用以加速学习
  • 在多种任务中实现跨环境对齐,提升泛化能力
  • 结果与大脑海马-内嗅系统相似,适合认知科学与AI研究者

人类和动物能以极少经验快速适应新环境,这依赖于情景记忆——即调用具体过往经历指导新情境决策。我们训练一个Transformer在一系列受啮齿类行为启发的规划任务中进行上下文内强化学习,并分析其内在学习机制。结果发现,表征学习依赖于上下文结构学习与跨情境对齐,使不同感官刺激下的表示得以统一。进一步表明,模型的强化学习策略无法用经典无模型或有模型规划解释,而是通过在记忆令牌中缓存中间计算,在决策时刻调用实现。整体上,记忆作为计算资源,同时存储原始经验与缓存结果,支持灵活行为。模型生成的表征与大脑海马-内嗅系统相关计算高度相似,提示该机制可能适用于自然认知研究。本工作为人工与自然系统中的快速适应提供了机制性假说。

原文摘要 · Abstract (English)

Humans and animals show remarkable learning efficiency, adapting to new environments with minimal experience. This capability is not well captured by standard reinforcement learning algorithms that rely on incremental value updates. Rapid adaptation likely depends on episodic memory -- the ability to retrieve specific past experiences to guide decisions in novel contexts. Transformers provide a useful setting for studying these questions because of their ability to learn rapidly in-context and because their key-value architecture resembles episodic memory systems in the brain. We train a transformer to in-context reinforcement learn in a distribution of planning tasks inspired by rodent behavior. We then characterize the learning algorithms that emerge in the model. We first find that representation learning is supported by in-context structure learning and cross-context alignment, where representations are aligned across environments with different sensory stimuli. We next demonstrate that the reinforcement learning strategies developed by the model are not interpretable as standard model-free or model-based planning. Instead, we show that in-context reinforcement learning is supported by caching intermediate computations within the model's memory tokens, which are then accessed at decision time. Overall, we find that memory may serve as a computational resource, storing both raw experience and cached computations to support flexible behavior. Furthermore, the representations developed in the model resemble computations associated with the hippocampal-entorhinal system in the brain, suggesting that our findings may be relevant for natural cognition. Taken together, our work offers a mechanistic hypothesis for the rapid adaptation that underlies in-context learning in artificial and natural settings.

强化学习记忆机制Transformer认知建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。