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

人类和动物在学习效率方面表现出色，能够凭借极少的经验便适应新的环境。标准的强化学习算法主要依赖于价值的逐步更新，难以很好地模拟这种能力。快速适应很可能依赖于情景记忆——即回忆特定过往经验以在新情境中做出决策的能力。由于Transformer具有在上下文中快速学习的能力，同时其键值架构类似于大脑中的情景记忆系统，因此为研究这些问题提供了一个有用的框架。 我们在一系列受啮齿类动物行为启发的规划任务中训练了一个Transformer模型，使其进行上下文内的强化学习，并分析模型中涌现出的学习机制。首先我们发现，表示学习得益于上下文内的结构学习以及跨上下文的对齐机制，即模型能够在感官刺激不同的环境中使表征保持一致。接着我们证明，该模型所发展的强化学习策略无法被解释为标准的无模型或基于模型的规划方法。相反，我们发现上下文内的强化学习依赖于模型内存标记中缓存的中间计算结果，并在决策时调用这些缓存信息。 总体而言，我们发现记忆不仅可以作为存储原始经验的资源，还可以存储中间计算结果，从而支持灵活的行为表现。此外，模型中发展出的表征与大脑海马-内嗅皮层系统的计算过程相似，表明我们的研究发现可能对理解自然认知具有启示意义。综合来看，本研究提出了一个机制性的假设，解释了人工和自然情境中上下文学习背后的快速适应能力。

论文旨在探讨人类和动物如何通过快速适应新环境（依赖于情景记忆）实现高效学习，而标准的强化学习算法往往依赖渐进式价值更新，难以捕捉这种能力。作者试图验证Transformer模型是否能够模拟这种快速适应机制，并揭示其背后的学习算法。

使用Transformer架构来模拟类脑的情景记忆系统，训练其在多种受啮齿类动物行为启发的任务中进行上下文内强化学习。关键在于发现模型并非采用传统模型无关或模型相关的规划策略，而是通过缓存中间计算并跨情境对齐表示结构来支持快速适应。

{提出一种新的机制假设：Transformer内存不仅存储原始经验，还缓存中间计算以支持灵活决策,实验设计模仿啮齿类动物的行为任务，具有生物学启发意义,发现表示学习与上下文结构学习、跨环境表示对齐密切相关,模型中的表示模式与大脑海马-内嗅皮层系统的计算相似，提示可能与自然认知机制相关,为理解人工和自然智能中的上下文学习提供了可解释的新视角}

{"Attention Is All You Need (Vaswani et al., 2017) - Transformer基础架构","Language Models as Knowledge Bases? (Petroni et al., 2019)","In-context Learning and Induction Heads (Elhage et al., 2021)","Neurobiology of intelligence: Science and society (Barrett et al., 2003)","The hippocampus as a cognitive graph (Hartley et al., 2014)","Model-Based Reinforcement Learning via Meta-Policy Optimization (Abdolmaleki et al., 2018)"}