- 奥地利林茨约翰开普勒大学机器学习研究所 ELLIS 小组，LIT 人工智能实验室
- 奥地利林茨 NXAI 有限公司
- 谷歌 DeepMind
- 米拉 - 魁北克人工智能研究所
摘要
近年来，强化学习（Reinforcement Learning, RL）领域出现了一种趋势，即通过序列建模在大规模数据集上离线训练大型动作模型。现有模型主要基于 Transformer 架构，这造就了强大的智能体。然而，由于推理时间较慢，基于 Transformer 的方法在实时应用（如机器人技术）中并不实用。最近，诸如 xLSTM 和 Mamba 等现代循环架构被提出，它们在训练期间展现出与 Transformer 架构类似的并行化优势，同时提供快速推理。在这项工作中，我们研究了这些现代循环架构用于大型动作模型的适用性。因此，我们提出了一种以 xLSTM 为核心的大型循环动作模型（Large Recurrent Action Model, LRAM），该模型具有线性时间推理复杂度和自然的序列长度外推能力。在来自 6 个领域的 432 个任务上进行的实验表明，LRAM 在性能和速度方面优于 Transformer。
1. 引言
强化学习（RL）已经取得了令人瞩目的成功，例如在游戏领域（Silver 等人，2016；Vinyals 等人，2019；Berner 等人，2019；Patil 等人，2022）、核聚变的等离子体控制（Degrave 等人，2022）以及平流层气球导航（Bellemare 等人，2020）。虽然这些成功基于经典的强化学习方法，即智能体通过 RL 目标进行在线训练，但最近出现了向离线 RL 设置（Levine 等人，2020；Schweighofer 等人，2022）和通过行为克隆训练的序列模型（Chen 等人，2021；Janner 等人，2021）发展的趋势。这种通过因果序列建模目标在大规模离线数据集上训练智能体的方法，受到基于 Transformer 架构的推动，并催生了我们所说的大型动作模型（Large Action Models, LAMs），以突出其与大型语言模型（LLMs）的相似性（Radford 等人，2018）。LAM 方法也可用于多任务设置，以开发通用智能体，如 Gato（Reed 等人，2022）。
现有的 LAM 主要基于 Transformer（Vaswani 等人，2017）架构。由于其强大的预测性能，机器人技术已成为大型模型的新兴应用领域（Brohan 等人，2023b；a；Octo Model Team 等人，2024；Gu 等人，2023；Wang 等人，2023），并且收集了大量多任务数据集（Jia 等人，2024；Embodiment Collaboration 等人，2024；Jiang 等人，2023；Mandlekar 等人，2023）。这一发展有望产生能够在广泛环境甚至不同实体中学习掌握复杂任务的机器人智能体。例如，最近已证明，尽管在有限的设置下，在多情节上下文中训练的序列模型可以执行上下文学习（ICL）（Laskin 等人，2020；Lee 等人，2023）。ICL 的一个潜在应用是在机器人技术中学习新的相关任务，而无需重新训练或微调。
In-Context Learning (ICL) 详解
1. 什么是 In-Context Learning (ICL)？
上下文学习（In-Context Learning, ICL） 指的是一种现象，其中模型可以通过观察输入提示（prompt）中的示例来学习新任务，而无需更新自身的参数（Brown et al., 2020）。这种能力在 GPT-3 及更大规模的预训练模型中被首次系统性研究，并被认为是大规模 Transformer 结构的重要特性（Vaswani et al., 2017; Radford et al., 2019）。
在 ICL 过程中，模型通过“示例-任务对”（example-task pairs）的形式，在一个输入序列（context window）内学习如何执行新任务。例如，给定几个翻译示例：
英：Hello, how are you? 法：Bonjour, comment ça va? 英：Good morning 法：模型可以自动推断翻译模式，并在不修改参数的情况下给出正确的翻译结果 "Bonjour"。
2. ICL 的工作原理
ICL 的核心在于 大规模自回归 Transformer 模型（Kaplan et al., 2020; Chan et al., 2022）的能力，即通过长序列的上下文窗口来“记住”或模式匹配任务规则，并直接在推理阶段执行新任务。它的主要机制包括：
- 模式匹配（Pattern Matching）：模型利用训练中积累的知识，在推理时匹配输入模式，以生成合理的输出（Garg et al., 2022）。
- 任务推理（Task Inference）：模型通过上下文示例推断出新任务的意图，而不依赖于显式的训练目标（Xie et al., 2022）。
- 隐式表示（Implicit Representations）：ICL 依赖于 Transformer 结构中的注意力机制，在高维表示空间中构建任务相关的隐式表示，使其能够在不同任务间进行泛化（Wei et al., 2023）。
3. ICL 在机器人领域的应用
机器人研究中，ICL 被认为是一种潜在的能力，可以使机器人学习新的任务，而不需要进行昂贵的微调（fine-tuning）或额外的数据收集。例如，近期的研究（Laskin et al., 2020; Lee et al., 2023）表明，基于 Transformer 的序列模型可以在多回合任务上下文中自动适应新的任务。
主要应用场景包括：
- 零样本任务适应（Zero-shot Task Adaptation）：机器人可以通过观察任务示例，在没有额外训练的情况下执行类似任务（Brohan et al., 2023）。
- 多模态信息整合（Multimodal Integration）：结合视觉、语言和运动数据，ICL 可用于机器人操作规划（Gu et al., 2023）。
- 跨环境任务迁移（Cross-Embodiment Generalization）：机器人可以在不同的环境或不同的硬件配置下执行任务，而无需重新训练（Octo Model Team et al., 2024）。
4. ICL 相关的重要论文
论文 贡献 Vaswani et al., 2017 提出了 Transformer 架构，奠定 ICL 机制的基础。 Brown et al., 2020 (GPT-3) 首次系统研究了 ICL 现象，并证明大模型可以在不微调的情况下进行任务泛化。 Laskin et al., 2020 在机器人任务中探索了 ICL，发现基于序列模型的机器人可以通过多轮交互适应新任务。 Lee et al., 2023 提出了一种基于 ICL 的机器人任务学习框架，在多个任务环境中取得了高泛化能力。 Gu et al., 2023 在机器人操作中应用 ICL，使机器人能够跨任务迁移并执行复杂操作。 Brohan et al., 2023 (RT-2) 提出了一种结合 ICL 和视觉-语言建模的机器人控制方法，使机器人能通过文本提示执行任务。 5. ICL 的挑战与未来研究方向
尽管 ICL 展现出强大的适应能力，但仍面临以下挑战：
- 任务泛化性限制：ICL 在部分任务上表现良好，但在复杂任务上仍然容易失败（Min et al., 2022）。
- 上下文窗口长度限制：Transformer 的固定窗口长度限制了 ICL 在长任务序列上的能力（Press et al., 2022）。
- 对示例质量的依赖：ICL 的成功与否很大程度上取决于提供的示例质量和顺序（Liu et al., 2023）。
未来的研究可能会围绕以下方向展开：
- 增强记忆能力：结合长时记忆（如 Mamba 模型，Gu et al., 2024）优化 ICL 在长时间序列上的表现。
- 自适应上下文长度：探索如 ALiBi（Press et al., 2022）等方法，提高 ICL 处理长序列任务的能力。
- 多模态 ICL：在机器人领域，结合视觉、语言和运动数据，使 ICL 适应更复杂的任务环境（Jia et al., 2024）。
ICL 是近年来大模型的重要能力之一，尤其在机器人领域有着广泛的应用潜力。通过观察上下文示例，模型可以在不更新参数的情况下适应新任务，提高泛化能力。虽然 ICL 在机器人控制、任务泛化和跨环境适应方面已有一定成果，但仍然面临任务复杂性、记忆能力和上下文窗口长度的挑战。未来研究可以结合 Mamba、Transformer 长序列建模、多模态融合等方法，进一步提升 ICL 在机器人任务中的表现。
基于 Transformer 的模型成功的关键原因之一是其通过训练期间的高效并行化扩展到大型数据集的能力。然而，尽管在 RL、语言建模（Brown 等人，2020）或计算机视觉（Dosovitskiy 等人，2021；He 等人，2022）等领域取得了众多成功案例，但基于 Transformer 的架构在速度和内存方面的推理成本一直较高（Kim 等人，2023）。因此，在资源受限的场景（如硬件容量有限和 / 或有实时约束的设备，如机器人或智能手机）中部署基于 Transformer 的模型是不可行的，因为需要快速的推理时间（Firoozi 等人，2023；Hu 等人，2023）。控制理论的一个基本原则是，控制器的采样率应与传感器的采样率处于同一数量级（Franklin 等人，1998，第 11 章）。例如，对于典型的机器人，如无人机或工业机械臂，为了保持系统稳定，需要 100Hz - 1000Hz 的采样率（Salzmann 等人，2023；El - Hussieny，2024；Hu 等人，2023；Chignoli 等人，2021），这意味着推理时间要小于 10 毫秒。在 1000Hz 的频率下，智能体 15 秒的运动对应于 15000 步的序列（El - Hussieny，2024），即使没有 ICL，也会导致较长的上下文长度。虽然存在一系列使大型模型更快的技术，如量化（Frantar 等人，2023）、蒸馏（Hinton 等人，2015）或剪枝（LeCun 等人，1989），但自注意力的二次时间复杂度仍然存在。
最近，现代循环架构被提出，它们在训练期间表现出与 Transformer 架构类似的并行化特性，同时具有线性时间推理复杂度。这些现代循环架构包括 xLSTM（Beck 等人，2024）和状态空间模型（State - Space Models, SSMs），如 Mamba（Gu 和 Dao，2023；Dao 和 Gu，2024）以及 Griffin/Hawk（De 等人，2024），它们不仅在语言建模领域挑战了 Transformer 的主导地位，还在计算机视觉（Alkin 等人，2024；Zhu 等人，2024）和生物医学（Schmidinger 等人，2024）等其他领域有所突破。更重要的是，它们的线性时间推理使其适合部署在计算资源有限、上下文规模大且有实时要求的场景中，如机器人技术。
在这项工作中，我们评估了现代循环架构（如 xLSTM 和 Mamba）作为大型动作模型的适用性。为此，我们引入了一种以 xLSTM 为核心的大型循环动作模型（LRAM）（见图 1）。我们使用类似于决策 Transformer（Decision Transformer, DT）（Chen 等人，2021）的监督学习设置，在来自 6 个领域的 432 个任务上训练我们的智能体。我们利用单任务专家智能体在线 RL 训练期间收集的数据，并将这些轨迹与其他专家演示一起编译成一个包含 8.94 亿次转换的大规模多领域数据集。由于其并行化特性，本文中考虑的现代循环架构可以像 Transformer 一样高效地处理这个大规模训练集，同时在推理时速度更快。对四种模型尺寸的多任务模型进行的实验表明，xLSTM 在性能和速度方面优于 Transformer。此外，我们研究了现代循环架构对微调性能和上下文学习能力的影响，发现它们在这两个方面都表现出色。
本文的主要目的是检验现代循环模型架构比 Transformer 更适合构建 LAM 的假设。在此过程中，我们做出了以下贡献：
- 我们提出了一种以 xLSTM 为核心的大型循环动作模型（LRAM），实现了高效推理。
- 我们评估了现代循环架构作为大型动作模型骨干网络在推理时间效率以及多任务、微调、上下文学习设置中的整体性能方面的适用性。
- 为了促进对大型动作模型的进一步研究，我们发布了数据准备管道和数据集。
2. 相关工作
2.1 RL 中的序列模型
长短期记忆网络（Long Short - Term Memory, LSTM）（Hochreiter 和 Schmidhuber，1997）是部分可观测在线 RL 问题的主要骨干架构，它助力实现了诸如掌握《星际争霸 II》（Vinyals 等人，2019）、《刀塔 2》（Berner 等人，2019）和雅达利游戏（Espeholt 等人，2018；Kapturowski 等人，2019）等成就。在 Transformer 在自然语言处理（NLP）（Devlin 等人，2019；Radford 等人，2019；Brown 等人，2020）、计算机视觉（Dosovitskiy 等人，2021；He 等人，2022；Radford 等人，2021；Fürst 等人，2022）和语音识别（Radford 等人，2022；Baevski 等人，2020）领域取得成功后，该架构也进入了 RL 领域。Chen 等人（2021）提出了决策 Transformer（DT），这是一种 GPT 风格的模型（Radford 等人，2018），通过行为克隆从离线轨迹中学习预测动作。轨迹 Transformer（Janner 等人，2021）预测动作以及状态和奖励，这使得动态建模成为可能。其他后续工作基于 DT（Zheng 等人，2022；Wang 等人，2022；Shang 等人，2022；Meng 等人，2021；Siebenborn 等人，2022；Schmied 等人，2024a），或者用 Mamba 替代 Transformer（Ota，2024；Dai 等人，2024）。此外，研究发现，在先前轨迹的条件下训练的序列模型表现出 ICL（Laskin 等人，2022；Lee 等人，2022；Kirsch 等人，2023），尽管是在有限的场景中。
2.2 大型动作模型（LAMs）
诸如决策 Transformer 之类的 LAM 非常适合多任务设置。Lee 等人（2022）发现，一个多游戏 DT 可以学习玩 46 种雅达利游戏。Reed 等人（2022）引入了一个在来自不同领域（从雅达利游戏到机械臂操作）的 600 多个任务上进行训练的通用智能体。Jiang 等人（2022）提出了一种基于多模态提示的用于机器人操作的 Transformer，它可以引导模型执行新任务。最近，Raad 等人（2024）引入了一个通过语言指令玩各种商业视频游戏的智能体。从那时起，由于大规模数据集的可用性（Jia 等人，2024；Embodiment Collaboration 等人，2024；Jiang 等人，2023；Mandlekar 等人，2023），机器人技术已成为开发 LAM 的新兴领域（Brohan 等人，2023b；a；Octo Model Team 等人，2024；Gu 等人，2023；Wang 等人，2023；Kim 等人，2024）。
2.3 下一代序列建模架构
线性循环模型，如状态空间模型（SSM，Gu 等人，2021；2022b；Smith 等人，2023；Orvieto 等人，2023），在长距离任务上挑战了 Transformer（Vaswani 等人，2017）架构的主导地位。这些线性循环神经网络（RNN）的关键见解是将循环状态矩阵对角化，并通过指数参数化实现稳定训练（Gu 等人，2022a；Orvieto 等人，2023）。从那时起，人们努力将 RNN 中的门控等特征融入其中（Elman，1990；Jordan，1990；Hochreiter 和 Schmidhuber，1997；Cho 等人，2014）。非线性门被认为具有更高的表达能力，但训练难度较大。Griffin（De 等人，2024）将门控线性循环与局部注意力相结合，比 Llama - 2（Touvron 等人，2023）实现了更高的训练数据效率和更好的序列外推。Mamba（Gu 和 Dao，2023）将一种类似于门控的选择机制引入到 SSM 中，使其状态和输入矩阵随时间变化。这与 RNN 的门控机制类似，但也与快速权重（Schmidhuber，1992）和线性注意力（Katharopoulos 等人，2020）等方法有相似之处。Mamba - 2（Dao 和 Gu，2024）强调了具有输入相关状态和输入矩阵的 SSM 与（门控）线性注意力变体之间的联系。最近，xLSTM（Beck 等人，2024）被提出，作为对经典 LSTM（Hochreiter 和 Schmidhuber，1997）的改进，它将门控、线性循环和循环权重组合在一个架构中用于语言建模。首先，xLSTM 利用带稳定化的指数门控为 RNN，更加强调重要输入。其次，xLSTM 由两个变体组成：mLSTM 变体侧重于记忆，在语言建模中被证明很重要；sLSTM 变体保留非对角化的循环矩阵以实现状态跟踪（Merrill 等人，2024）。状态跟踪在逻辑任务中很重要，而像 Mamba、Griffin 或 Transformer 这样的线性化循环或状态空间模型从根本上无法对其进行建模。
3. 大型循环动作模型
3.1 背景
强化学习：
我们假设通过马尔可夫决策过程（Markov Decision Process, MDP）来表示标准的 RL 公式，用元组
表示，其中和分别表示状态空间和动作空间。在每个时间步，智能体观察到状态，预测动作，并获得一个标量奖励。奖励由奖励函数确定，定义了转移动态，表示在状态执行动作时，下一个状态的概率分布。RL 的目标是学习一个策略，该策略在状态预测动作，以最大化。
决策 Transformer Decision Transformer（Chen 等人，2021）
通过因果序列建模将 RL 问题设置转化为下一个动作预测任务。在训练时，DT 旨在学习一个策略，将未来奖励映射到动作，这通常被称为倒置 RL upside-down RL（Schmidhuber，2019）。在推理时，DT 通过目标回报进行条件设定，以发出高奖励动作。因此，我们假设可以访问一个数据集
，其中包含个轨迹，每个轨迹由状态、未来回报（return - to - go, RTG）、动作和奖励组成的四元组
构成。这里，T 表示轨迹的长度。DT 被训练以根据数据集中的子轨迹预测真实动作：
其中是上下文窗口的大小。实际上，公式（1）描述了多游戏 DT（Lee 等人，2022）的设置，它在序列表示中也包括了奖励。
3.2 大型循环动作模型（LRAMs）
我们的 LRAM 以现代循环架构为核心（见图 1），具有并行训练和循环推理模式。我们用三种不同的变体实例化 LRAM，即两种不同的 xLSTM 配置和 Mamba。我们采用与 Lee 等人（2022）和 Reed 等人（2022）类似的训练协议，但有一些重要差异，旨在加快不同骨干网络的推理速度。
多模态序列表示：为了编码来自不同环境、具有不同状态和动作空间的输入，我们为每个模态使用单独的编码器，这些编码器在任务和领域之间共享。对于编码图像，我们使用类似于 Espeholt 等人（2018）的卷积神经网络（Convolutional Neural Network, CNN），而对于低维输入，我们使用全连接网络。我们避免对图像进行分块和对连续状态进行标记化，以避免不必要的长序列。同样，我们使用线性层对奖励和 RTG 进行编码。在我们的序列公式中，我们省略了动作，因为我们发现这可能会对性能产生不利影响，特别是对于动作变化平滑的连续控制任务（见 4.3 节）。因此，我们的轨迹形式为
，并且我们训练策略
来预测真实动作
，如下所示：
共享动作头：强化学习中的动作空间通常因环境而异。例如，在我们考虑的环境中，有 18 个离散动作，对于连续控制环境，最多有 8 个连续维度。因此，我们将连续动作维度离散化为 256 个均匀间隔的区间，类似于 Reed 等人（2022 年）和 Brohan 等人（2023b）的做法。与先前的工作不同，我们利用一个共享的动作头来联合预测所有离散动作或连续动作维度。我们发现，与使用连续动作的自回归动作预测相比，这种设置显著减少了推理时间。
循环推理模式：
在推理时，我们利用循环骨干网络并保留上一个时间步的隐藏状态。这使得沿着序列长度的推理具有线性时间复杂度，从而实现快速推理。此外，循环式推理非常适合通过强化学习目标进行在线微调，类似于在线强化学习中基于长短期记忆网络的策略。为了加快推理速度，我们为 xLSTM 骨干网络使用了定制内核（见附录 21）。
我们统一的离散动作表示使得我们的智能体能够通过交叉熵损失作为所有任务和领域的训练目标进行一致的训练，类似于 Reed 等人（2022 年）的方法。我们为每个领域使用单独的奖励尺度，并为每个任务设置目标回报。此外，我们没有使用 Chen 等人（2021 年）所使用的时间步编码，因为当情节长度变化时，这种编码是有害的。我们在附录 C 中提供了更多的实现细节。
4. 实验
我们研究了现代循环架构作为大型动作模型在来自 6 个领域（雅达利（Bellemare 等人，2013 年）、Composuite（Mendez 等人，2022 年）、DMControl（Tassa 等人，2018 年）、Meta-World（Yu 等人，2020b）、Mimicgen（Mandlekar 等人，2023 年）和 Procgen（Cobbe 等人，2020b））的 432 个任务上的适用性。为此，我们编译了一个包含 8.94 亿次转换的大规模数据集（见 4.1 节）。在所有实验中，我们比较了四种骨干网络变体：xLSTM [7:1]、xLSTM [1:0]（Beck 等人，2024 年）、Mamba（Gu 和 Dao，2023 年），以及决策 Transformer（Chen 等人，2021 年）中使用的 GPT-2 风格的 Transformer。按照 Beck 等人（2024 年）的方法，我们对 xLSTM 使用括号表示法，该表示法表示 mLSTM 与 sLSTM 块的比例。例如，xLSTM [1:0] 仅包含 mLSTM 块。
在 4.2 节中，我们对从 1600 万到 2.08 亿 16M to 208M 参数的四种模型尺寸进行了缩放比较，结果表明，现代循环架构在不同模型尺寸下的性能与 Transformer 基线相当或更优。在 4.3 节中，我们研究了循环骨干网络对微调性能、上下文学习能力的影响，并进一步分析了我们训练的循环骨干网络。最后，在 4.4 节中，我们通过实验检验了基于 xLSTM 和基于 Transformer 的智能体在推理时间方面的差异，包括延迟和吞吐量，这表明循环骨干网络具有优势。
4.1 数据集和环境
数据集：我们编译了一个包含来自六个领域的 432 个任务的大规模数据集。如果有先前工作的数据集，我们会加以利用。对于雅达利游戏，我们从 Agarwal 等人（2020 年）发布的 DQN-Replay 数据集中为每个任务提取 500 万个转换。对于 Composuite，我们利用（Hussing 等人，2023 年）发布的数据集。对于 Meta-World，我们使用（Schmied 等人，2024a）发布的每个任务 200 万个转换。对于 DMControl，我们使用特定任务的强化学习智能体为每个任务生成 1000 万个转换。对于 Mimicgen，我们使用（Mandlekar 等人，2023 年）发布的 21 个任务的数据集，并为其余 62 个任务生成轨迹。最后，对于 Procgen，我们从（Schmied 等人，2024b）发布的数据集中提取 2000 万个转换。我们的最终数据集包含 340 万个轨迹，总共 8.94 亿次转换（见表 2）。我们从相同领域中保留了另外 37 个任务用于零样本评估。为了促进未来的研究，我们发布了我们的数据准备管道和生成的数据。
环境：雅达利和 Procgen 提供图像观察和离散动作。相比之下，其余四个领域呈现基于状态的观察和连续动作。因此，我们的实验涉及状态和动作空间的混合，以及不同的情节长度（见表 2）。按顺序在所有 432 个任务上定期评估训练好的智能体非常耗时，因此我们将评估分布在 GPU 和并行进程中（见附录 C）。关于我们的数据集和环境的更多详细信息，请参见附录 B。
4.2 缩放比较
为了进行主要比较，我们在 432 个任务的完整训练任务混合集上训练我们的四种骨干网络变体。对于每个架构骨干网络，我们报告四种模型尺寸（1600 万、4800 万、1.08 亿和 2.06 亿参数）的性能得分。我们使用 128 的批量大小和 50 个时间步的上下文长度对所有模型进行 20 万次更新训练。所有领域的比例大致相等，每个领域导致 3.3 万次更新。附录 C 中提供了每个骨干网络变体和模型尺寸的更多实现细节和超参数。
序列预测性能：在图 2a 中，我们报告了所有骨干网络和模型尺寸在各个领域的单独得分上平均的验证集困惑度。为此，我们为每个训练任务保留一组轨迹（2.5%），并在每 5 万步后计算困惑度（训练困惑度见图 12）。两种循环骨干网络在很大程度上优于 Transformer 基线，特别是随着模型尺寸的增加。
图2. 缩放比较。我们比较了xLSTM、Mamba和决策Transformer（DT）在四种模型规模下的表现：参数分别为1600万、4800万、1.1亿和2.06亿。我们展示了（a）在留出数据集上的验证困惑度，以及（b）在训练任务环境中评估得到的归一化分数，该分数是在所有6个领域上取平均值。
评估性能：在训练过程中，我们每 5 万步在所有 432 个训练环境中评估我们的智能体。在图 2b 中，我们报告了在所有六个领域上平均的归一化性能结果。循环骨干网络在所有模型尺寸上都优于 Transformer。虽然 xLSTM 和 Mamba 在较小规模下性能相似，但 xLSTM 在较大规模（2.06 亿）下往往优于 Mamba。这是 xLSTM 的一个重要优势，因为 LRAM 智能体可以从更多数据和更大模型中受益匪浅。请注意，Mamba 的参数数量明显多于竞争对手。关于在 37 个保留任务上的零样本评估性能，请参见附录 D.2 中的图 14。
每个领域的性能：在图 3 中，我们报告了 2.06 亿模型在所有六个领域上获得的归一化分数。对于 Meta-World、DMControl、Mimicgen、Composuite 和 Procgen，我们使用（Levine 等人，2020 年）建议的数据归一化分数。对于雅达利，我们报告人类归一化分数。我们观察到，xLSTM 骨干网络在六个领域中的三个领域上优于竞争对手，而在其余领域上性能相似。、
4.3 分析与消融实验
微调：为了评估循环骨干网络对微调性能的影响，我们在来自所有 6 个领域的 37 个保留环境上对模型进行微调。我们评估了 1600 万参数预训练模型的 xLSTM 架构的微调性能，并将其与从头开始训练的 xLSTM 进行比较。预训练的 LRAM 在大多数领域中优于随机初始化的 xLSTM 模型（见图 15）。这表明切换骨干网络不会对微调性能产生负面影响。
上下文学习：接下来，我们在先前关于上下文强化学习的工作（Laskin 等人，2022 年；Lee 等人，2023 年；Schmied 等人，2024b）中考虑的暗室环境中研究循环骨干网络的上下文学习能力。为了单独研究上下文学习，我们使用多情节上下文从头开始训练模型，这导致了较长的上下文长度（实验设置的详细信息见附录 D.4）。特别是，我们采用算法蒸馏（AD，Laskin 等人，2022 年）框架，并将 Transformer 骨干网络架构替换为现代循环架构。在图 4 中，我们报告了在 20 个保留任务上的上下文学习性能（训练任务见图 16）。我们发现 xLSTM [7:1] 在 80 个训练任务和 20 个保留任务上都获得了最高的总体分数，我们将其归因于 sLSTM 块的状态跟踪能力（Merrill 等人，2024 年）。
嵌入空间分析：在图 5 中，我们分析了我们模型学习到的表示。我们从每个任务中采样 32 个子轨迹，提取最后一层的序列表示，使用 UMAP（McInnes 等人，2018 年）对它们进行聚类，并根据其领域对每个点进行着色（更多详细信息见附录 F）。我们发现来自同一领域的任务聚集在一起。此外，与 DT 相比，xLSTM 表现出更精细的领域分离，这可能进一步有助于更好的下游性能。
去除动作和上下文长度的影响：我们发现从上下文中去除动作会在所有骨干网络上带来更好的性能。虽然上下文长度超过 1 会损害 Meta-World 和 DMControl 上的性能，并且在训练中包含动作时也是如此，但在不包含动作进行训练时，情况则相反（见图 22、23、25）。这与最近的工作形成对比，这些工作没有从更长的收集中受益（Octo Model Team 等人，2024 年）。虽然去除动作提高了 Meta-World/DMControl 上的性能，但它对离散控制环境的性能没有影响。在 Meta-World/DMControl 上，我们观察到模型变得过于自信，如果产生不良的初始动作，这将是一个问题。这是因为许多机器人环境中的动作变化平滑，通过观察先前的动作，智能体学习到了捷径。Wen 等人（2020 年）也观察到了类似的问题，并将其称为模仿问题。从输入中去除动作可以防止智能体使用捷径，并减轻模仿问题。重要的是，随着序列长度的增加，跨领域的评估性能会提高，这表明历史有助于预测下一个动作（例如，通过观察过去犯的错误，见图 24、26）。
回报条件与行为克隆
Return-conditioning vs. Behavior Cloning：
在我们的实验中，除了上下文学习实验外，我们使用了一种包含未来回报令牌 return-to-go token（RTG）的序列表示，这在 DT 文献（Chen 等人，2021 年；Lee 等人，2022 年）中很常见。在推理时，RTG 允许模型根据高目标回报进行条件设定，以产生高质量的动作。这在数据集中包含最优和次优轨迹的混合时特别有用。然而，最近的许多工作专注于没有回报条件的行为克隆（Reed 等人，2022 年；Brohan 等人，2023a）。因此，我们研究了在 2.06 亿参数规模下从序列中排除 RTG / 奖励令牌的效果，以验证我们的发现是否适用于行为克隆设置。确实，我们发现相同的趋势仍然成立（见图 27 和 28）。
图27. 省略未来回报（RTG）条件的影响的消融实验。我们报告了2.06亿参数模型在432个训练任务上的（a）验证困惑度和（b）评估性能的学习曲线。我们观察到与在序列中包含RTG时相似的性能趋势。
图28. 省略未来回报（RTG）条件和奖励条件的影响的消融实验。我们报告了2.06亿参数模型在432个训练任务上的（a）验证困惑度和（b）评估性能的学习曲线。我们观察到与在序列中包含RTG时相似的性能趋势。
mLSTM 与 sLSTM 的比例：
在整个实验中，我们比较了两种 xLSTM 变体：xLSTM [7:1] 和 xLSTM [1:0]。这种括号表示法由 Beck 等人（2024 年）提出，为了保持一致性，我们保持相同的比例（见附录 C.3）。虽然 mLSTM 是可并行化的，但 sLSTM 实现了状态跟踪（Merrill 等人，2024 年）。为了更好地理解该比例的影响，我们在 432 个任务和暗室环境上进行了消融实验（见附录 E.3），类似于 Beck 等人（2024 年）的方法。我们发现其他比例，如 [3:1]，也可能是有效的，并强调了在较低层放置 sLSTM 的重要性（图 30）。然而，sLSTM 层的有效性取决于手头的任务。具有长视野或部分可观测性的复杂任务，如现实世界应用中常见的任务，可能会从 sLSTM 提供的状态跟踪能力中受益。
我们在附录 E.5 和 E.4 中分别提供了关于减少 xLSTM 层数和禁用 DT 中 Dropout 的影响的更多消融实验。
4.4 推理时间比较
最后，我们通过实验检验了基于 xLSTM 和基于 Transformer 的智能体在推理时间方面的差异。与 De 等人（2024 年）类似，我们报告了延迟和吞吐量。由于延迟对于实时应用更为重要，因此我们主要分析延迟。
设置：
我们在具有 40GB RAM 的 A100 上使用 2.06 亿参数的模型进行所有推理时间测试。对于 Transformer，我们使用 PyTorch（Paszke 等人，2019 年）支持的 KV 缓存和 FlashAttention（Dao，2023 年）。对于 xLSTM，我们使用自定义内核的循环式推理来加速计算（内核加速的影响见图 21）。对于两种骨干网络，我们都使用 torch.compile。具有 KV 缓存的 Transformer 每步具有线性时间复杂度，在序列长度上具有二次复杂度。相比之下，xLSTM 每步具有恒定时间复杂度，在序列长度上具有线性复杂度。因此，正如 De 等人（2024 年）所观察到的，我们预计在更长的序列和更大的批量大小下会有加速。为了确保公平比较，我们比较具有相同层数块的 DT 和 xLSTM，并增加 xLSTM 的隐藏大小以匹配 DT 的参数数量（这些模型的评估性能见附录 E.5）。我们在附录 D.5 中提供了推理时间测试的更多详细信息。
环境：
我们在实验中平均情节长度最长的环境 —— 雅达利游戏《高速公路》（Freeway）上进行所有推理时间测试。《高速公路》中的每个情节持续 8192 步，相当于 24576 个令牌（状态 / RTG / 奖励）。我们对所有模型进行 5 个情节的评估，并在情节边界之间保留 KV 缓存 / 隐藏状态。报告的延迟和吞吐量是在所有评估情节上的平均值，但不包括第一个情节，我们排除第一个情节是为了排除编译时间和预填充。我们选择在环境交互期间测量推理时间，即包括模拟器延迟，而不仅仅是令牌生成。
延迟：
与 De 等人（2024 年）类似，我们通过固定批量大小 B 执行单个推理步骤的平均时间（秒）来测量延迟（越低越好）。在图 6 中，我们报告了不同上下文长度和两个批量大小的延迟。请注意，是以时间步为单位，每个时间步包含 3 个令牌（状态、未来回报、奖励）。因此，最大的有效序列长度为 76800。正如预期的那样，我们发现循环骨干网络的推理延迟低于 Transformer。随着序列长度的增加，由于 KV 缓存大小的增加，DT 会耗尽内存（见图 6c）。相比之下，xLSTM 的推理速度与上下文长度无关，因此可以实现明显更长的上下文长度。这种特性对于需要在上下文中保留多个情节的上下文强化学习特别有意义（Laskin 等人，2022 年）。然而，我们的实验强调，复杂度优势的实现取决于设备、模型大小、批量大小和上下文长度，这与 De 等人（2024 年）的发现类似。
图6. 在A100上的延迟比较。我们报告了不同上下文长度（以时间步为单位）下的延迟情况，其中批量大小分别为（a）$B = 1$和（b）$B = 16$ 。在（c）中，我们展示了批量大小$B = 1$时，GPU内存的消耗百分比。我们在雅达利《高速公路》游戏环境中，对具有相同层数块和相同参数数量的决策Transformer（DT）和xLSTM进行比较。DT缺失的条形表示内存不足（OOM）。
吞吐量：吞吐量通过具有固定上下文长度的模型每秒执行的推理步骤总数来衡量。在图 7 中，我们报告了固定上下文长度时不同批量大小的吞吐量。这里，批量大小可以解释为智能体与之交互的并行环境的数量。正如预期的那样，我们发现 xLSTM 的吞吐量明显高于 DT。xLSTM 的优势随着批量大小的增加而增加。虽然在序列长度上具有二次复杂度的 DT 在批量大小超过 64 时会耗尽内存，但具有线性复杂度的 xLSTM 可以轻松处理更大的批量大小。在这两个实验中，循环 xLSTM 的性能优于 Transformer 骨干网络。
图7. 在A100上，针对雅达利《高速公路》环境，当上下文长度$C = 1600$ 时间步时，不同批量大小下的吞吐量比较。决策Transformer（DT）缺失的条形表示内存不足（OOM）。
5. 结论
在这项工作中，我们研究了现代循环架构作为 Transformer 的替代方案，用于构建大型动作模型（LAMs）的适用性。我们发现，以 xLSTM 或 Mamba 为核心的 LRAM，在不同模型规模下的评估性能优于 Transformer（见 4.2 节）。此外，我们证明了基于 xLSTM 的 LRAM 推理速度更快，尤其是在上下文规模较大的情况下（见 4.4 节）。因此，实证证据表明，循环架构的骨干网络对于 LAMs 而言，是具有吸引力的替代选择。值得注意的是，xLSTM 的线性时间推理复杂度，可能使需要长上下文长度的应用（如上下文学习）成为可能，并有助于将大规模智能体应用于机器人技术等实时应用场景。
现代循环架构和 Transformer 各有优缺点。一方面，xLSTM 和 Mamba 相较于 Transformer 具有复杂度优势。它们的线性复杂度确保了计算需求随序列长度的增加而增长得更慢，从而实现更高效的推理，这对于边缘应用尤为重要。虽然我们在高端数据中心 GPU 上进行推理时间的比较，但边缘设备上的应用可能需要应对性能较弱的加速器。重要的是，我们发现 LAMs 能从更长的序列中显著受益（见 4.3 节）。另一方面，Transformer 对于需要精确回忆序列中标记的应用非常有效，这在决策过程中可能至关重要（Ni 等人，2024）。最后，xLSTM 尤其能够通过 sLSTM 块实现状态跟踪，而 Transformer 和 Mamba 无法做到这一点（Merrill 等人，2024）。状态跟踪在逻辑任务或处理部分可观测性时非常重要，对于实际应用的从业者来说，这可能是一个有用的工具。鉴于这些差异，应根据具体任务考虑选择不同的骨干网络。
局限性：LAMs 的主要目标应用是机器人技术。虽然我们的大多数实验涉及机器人模拟，但尚未在实际机器人上进行实验。然而，我们相信我们的研究结果可以推广到现实场景，并计划在未来的工作中提供更多证据。此外，我们的微调实验仅限于离线 RL。我们设想，在大规模数据集上预训练的智能体可以通过在线 RL 进行成功微调，以探索训练数据中未出现的新策略。现代循环架构同时提供并行和循环训练模式，这可能是此类应用成功的关键。虽然我们证明了 LRAM 的上下文学习能力有所提升，但仅考虑了网格世界的设置。我们旨在进一步研究 LRAM 在更复杂环境中的上下文学习能力。
